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 { AIMessage } from "@langchain/core/messages";
import { createMiddleware } from "langchain";
function getTurnKey(context) {
    return `${context.sessionId}:${context.turnId}`;
}
function getResponseId(message) {
    var _a;
    const metadata = message.response_metadata;
    return (_a = metadata === null || metadata === void 0 ? void 0 : metadata.id) !== null && _a !== void 0 ? _a : null;
}
function getPreviousResponseId(modelSettings) {
    return modelSettings === null || modelSettings === void 0 ? void 0 : modelSettings.previous_response_id;
}
function getContinuationMessages(messages, previousResponseId) {
    var _a;
    let continuationStartIndex = null;
    for (let index = messages.length - 1; index >= 0; index -= 1) {
        const message = messages[index];
        if (AIMessage.isInstance(message) &&
            ((_a = message.response_metadata) === null || _a === void 0 ? void 0 : _a.id) ===
                previousResponseId) {
            continuationStartIndex = index + 1;
            break;
        }
    }
    if (continuationStartIndex === null) {
        return null;
    }
    return messages.slice(continuationStartIndex);
}
export function createOpenAiResponsesContinuationMiddleware() {
    const responseIdsByTurn = new Map();
    return createMiddleware({
        name: "OpenAiResponsesContinuationMiddleware",
        wrapModelCall(request, handler) {
            return __awaiter(this, void 0, void 0, function* () {
                var _a;
                const context = request.runtime.context;
                const turnKey = getTurnKey(context);
                const previousResponseId = (_a = getPreviousResponseId(request.modelSettings)) !== null && _a !== void 0 ? _a : responseIdsByTurn.get(turnKey);
                const continuationMessages = previousResponseId
                    ? getContinuationMessages(request.messages, previousResponseId)
                    : null;
                const response = yield handler(previousResponseId && continuationMessages
                    ? Object.assign(Object.assign({}, request), { messages: continuationMessages, modelSettings: Object.assign(Object.assign({}, request.modelSettings), { previous_response_id: previousResponseId }) }) : request);
                const responseId = getResponseId(response);
                if (responseId) {
                    responseIdsByTurn.set(turnKey, responseId);
                }
                else {
                    responseIdsByTurn.delete(turnKey);
                }
                return response;
            });
        },
    });
}
