"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.LangSmithMiddleware = LangSmithMiddleware;
const traceable_js_1 = require("../../traceable.cjs");
const vercel_js_1 = require("../../utils/vercel.cjs");
const utils_js_1 = require("./utils.cjs");
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const _formatTracedInputs = (params) => {
    const { prompt, ...rest } = params;
    if (prompt == null) {
        return params;
    }
    if (Array.isArray(prompt)) {
        return {
            ...rest,
            messages: prompt.map((message) => (0, utils_js_1.convertMessageToTracedFormat)(message)),
        };
    }
    return rest;
};
const _formatTracedOutputs = (outputs) => {
    const formattedOutputs = { ...outputs };
    if (formattedOutputs.role == null) {
        formattedOutputs.role = formattedOutputs.type ?? "assistant";
    }
    return (0, utils_js_1.convertMessageToTracedFormat)(formattedOutputs);
};
const setUsageMetadataOnRunTree = (result, runTree) => {
    if (result.usage == null || typeof result.usage !== "object") {
        return;
    }
    // Shim for AI SDK 4
    const inputTokens = result.usage?.inputTokens ??
        result.usage?.promptTokens;
    const outputTokens = result.usage?.outputTokens ??
        result.usage?.completionTokens;
    let totalTokens = result.usage?.totalTokens;
    if (typeof totalTokens !== "number" &&
        typeof inputTokens === "number" &&
        typeof outputTokens === "number") {
        totalTokens = inputTokens + outputTokens;
    }
    const langsmithUsage = {
        input_tokens: inputTokens,
        output_tokens: outputTokens,
        total_tokens: totalTokens,
    };
    const inputTokenDetails = (0, vercel_js_1.extractInputTokenDetails)(result.usage, result.providerMetadata);
    const outputTokenDetails = (0, vercel_js_1.extractOutputTokenDetails)(result.usage, result.providerMetadata);
    runTree.extra = {
        ...runTree.extra,
        metadata: {
            ...runTree.extra?.metadata,
            usage_metadata: {
                ...langsmithUsage,
                input_token_details: {
                    ...inputTokenDetails,
                },
                output_token_details: {
                    ...outputTokenDetails,
                },
            },
        },
    };
};
/**
 * AI SDK middleware that wraps an AI SDK 5 model and adds LangSmith tracing.
 */
function LangSmithMiddleware(config) {
    const { name, modelId, lsConfig } = config ?? {};
    return {
        wrapGenerate: async ({ doGenerate, params }) => {
            const traceableFunc = (0, traceable_js_1.traceable)(async (_params) => {
                const result = await doGenerate();
                const currentRunTree = (0, traceable_js_1.getCurrentRunTree)(true);
                if (currentRunTree !== undefined) {
                    setUsageMetadataOnRunTree(result, currentRunTree);
                }
                return result;
            }, {
                ...lsConfig,
                name: name ?? "ai.doGenerate",
                run_type: "llm",
                metadata: {
                    ls_model_name: modelId,
                    ai_sdk_method: "ai.doGenerate",
                    ...lsConfig?.metadata,
                },
                processInputs: (inputs) => {
                    const typedInputs = inputs;
                    const inputFormatter = lsConfig?.processInputs ?? _formatTracedInputs;
                    return inputFormatter(typedInputs);
                },
                processOutputs: (outputs) => {
                    const typedOutputs = outputs;
                    const outputFormatter = lsConfig?.processOutputs ?? _formatTracedOutputs;
                    return outputFormatter(typedOutputs);
                },
            });
            const res = await traceableFunc(params);
            return res;
        },
        wrapStream: async ({ doStream, params }) => {
            const parentRunTree = (0, traceable_js_1.getCurrentRunTree)(true);
            let runTree;
            if (parentRunTree != null &&
                typeof parentRunTree === "object" &&
                typeof parentRunTree.createChild === "function") {
                const inputFormatter = lsConfig?.processInputs ?? _formatTracedInputs;
                const formattedInputs = inputFormatter(params);
                runTree = parentRunTree?.createChild({
                    ...lsConfig,
                    name: name ?? "ai.doStream",
                    run_type: "llm",
                    metadata: {
                        ls_model_name: modelId,
                        ai_sdk_method: "ai.doStream",
                        ...lsConfig?.metadata,
                    },
                    inputs: formattedInputs,
                });
            }
            await runTree?.postRun();
            try {
                const { stream, ...rest } = await doStream();
                const chunks = [];
                const transformStream = new TransformStream({
                    async transform(chunk, controller) {
                        if (chunk.type === "tool-input-start" ||
                            chunk.type === "text-start") {
                            // Only necessary to log the first token event
                            if (runTree?.events == null ||
                                (Array.isArray(runTree.events) && runTree.events.length === 0)) {
                                runTree?.addEvent({ name: "new_token" });
                            }
                        }
                        else if (chunk.type === "finish") {
                            runTree?.addEvent({ name: "end" });
                        }
                        chunks.push(chunk);
                        controller.enqueue(chunk);
                    },
                    async flush() {
                        try {
                            const output = chunks.reduce((aggregated, chunk) => {
                                if (chunk.type === "text-delta") {
                                    if (chunk.delta != null) {
                                        return {
                                            ...aggregated,
                                            content: aggregated.content + chunk.delta,
                                        };
                                    }
                                    else if ("textDelta" in chunk &&
                                        chunk.textDelta != null) {
                                        // AI SDK 4 shim
                                        return {
                                            ...aggregated,
                                            content: aggregated.content + chunk.textDelta,
                                        };
                                    }
                                    else {
                                        return aggregated;
                                    }
                                }
                                else if (chunk.type === "tool-call") {
                                    const matchingToolCall = aggregated.tool_calls.find((call) => call.id === chunk.toolCallId);
                                    if (matchingToolCall != null) {
                                        return aggregated;
                                    }
                                    let chunkArgs = chunk.input;
                                    if (chunkArgs == null &&
                                        "args" in chunk &&
                                        typeof chunk.args === "string") {
                                        chunkArgs = chunk.args;
                                    }
                                    return {
                                        ...aggregated,
                                        tool_calls: [
                                            ...aggregated.tool_calls,
                                            {
                                                id: chunk.toolCallId,
                                                type: "function",
                                                function: {
                                                    name: chunk.toolName,
                                                    arguments: chunkArgs,
                                                },
                                            },
                                        ],
                                    };
                                }
                                else if (chunk.type === "finish") {
                                    if (runTree != null) {
                                        setUsageMetadataOnRunTree(chunk, runTree);
                                    }
                                    return {
                                        ...aggregated,
                                        providerMetadata: chunk.providerMetadata,
                                        finishReason: chunk.finishReason,
                                    };
                                }
                                else {
                                    return aggregated;
                                }
                            }, {
                                content: "",
                                role: "assistant",
                                tool_calls: [],
                            });
                            const outputFormatter = lsConfig?.processOutputs ?? utils_js_1.convertMessageToTracedFormat;
                            const formattedOutputs = await outputFormatter(output);
                            await runTree?.end(formattedOutputs);
                        }
                        catch (error) {
                            await runTree?.end(undefined, error.message ?? String(error));
                            throw error;
                        }
                        finally {
                            await runTree?.patchRun({
                                excludeInputs: true,
                            });
                        }
                    },
                });
                return {
                    stream: stream.pipeThrough(transformStream),
                    ...rest,
                };
            }
            catch (error) {
                await runTree?.end(undefined, error.message ?? String(error));
                await runTree?.patchRun({
                    excludeInputs: true,
                });
                throw error;
            }
        },
    };
}
