"use strict";
/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation and GitHub. All rights reserved.
 *--------------------------------------------------------------------------------------------*/
Object.defineProperty(exports, "__esModule", { value: true });
exports.LogicalWrapper = exports.IfEmpty = exports.AbstractKeepWith = exports.TokenLimit = exports.Expandable = exports.Chunk = exports.LegacyPrioritization = exports.ToolResult = exports.PrioritizedList = exports.Document = exports.Image = exports.TextChunk = exports.ToolMessage = exports.AssistantMessage = exports.UserMessage = exports.SystemMessage = exports.BaseChatMessage = void 0;
exports.isChatMessagePromptElement = isChatMessagePromptElement;
exports.useKeepWith = useKeepWith;
const _1 = require(".");
const promptElement_1 = require("./promptElement");
function isChatMessagePromptElement(element) {
    return (element instanceof SystemMessage ||
        element instanceof UserMessage ||
        element instanceof AssistantMessage);
}
class BaseChatMessage extends promptElement_1.PromptElement {
    render() {
        return vscpp(vscppf, null, this.props.children);
    }
}
exports.BaseChatMessage = BaseChatMessage;
/**
 * A {@link PromptElement} which can be rendered to an OpenAI system chat message.
 *
 * See {@link https://platform.openai.com/docs/api-reference/chat/create#chat-create-messages}
 */
class SystemMessage extends BaseChatMessage {
    constructor(props) {
        props.role = _1.Raw.ChatRole.System;
        super(props);
    }
}
exports.SystemMessage = SystemMessage;
/**
 * A {@link PromptElement} which can be rendered to an OpenAI user chat message.
 *
 * See {@link https://platform.openai.com/docs/api-reference/chat/create#chat-create-messages}
 */
class UserMessage extends BaseChatMessage {
    constructor(props) {
        props.role = _1.Raw.ChatRole.User;
        super(props);
    }
}
exports.UserMessage = UserMessage;
/**
 * A {@link PromptElement} which can be rendered to an OpenAI assistant chat message.
 *
 * See {@link https://platform.openai.com/docs/api-reference/chat/create#chat-create-messages}
 */
class AssistantMessage extends BaseChatMessage {
    constructor(props) {
        props.role = _1.Raw.ChatRole.Assistant;
        super(props);
    }
}
exports.AssistantMessage = AssistantMessage;
const WHITESPACE_RE = /\s+/g;
/**
 * A {@link PromptElement} which can be rendered to an OpenAI tool chat message.
 *
 * See {@link https://platform.openai.com/docs/api-reference/chat/create#chat-create-messages}
 */
class ToolMessage extends BaseChatMessage {
    constructor(props) {
        props.role = _1.Raw.ChatRole.Tool;
        super(props);
    }
}
exports.ToolMessage = ToolMessage;
/**
 * A chunk of single-line or multi-line text that is a direct child of a {@link ChatMessagePromptElement}.
 *
 * TextChunks can only have text literals or intrinsic attributes as children.
 * It supports truncating text to fix the token budget if passed a {@link TextChunkProps.tokenizer} and {@link TextChunkProps.breakOn} behavior.
 * Like other {@link PromptElement}s, it can specify `priority` to determine how it should be prioritized.
 */
class TextChunk extends promptElement_1.PromptElement {
    async prepare(sizing, _progress, token) {
        const breakOn = this.props.breakOnWhitespace ? WHITESPACE_RE : this.props.breakOn;
        if (!breakOn) {
            return vscpp(vscppf, null, this.props.children);
        }
        let fullText = '';
        const intrinsics = [];
        for (const child of this.props.children || []) {
            if (child && typeof child === 'object') {
                if (typeof child.ctor !== 'string') {
                    throw new Error('TextChunk children must be text literals or intrinsic attributes.');
                }
                else if (child.ctor === 'br') {
                    fullText += '\n';
                }
                else {
                    intrinsics.push(child);
                }
            }
            else if (child != null) {
                fullText += child;
            }
        }
        const text = await getTextContentBelowBudget(sizing, breakOn, fullText, token);
        return (vscpp(vscppf, null,
            intrinsics,
            text));
    }
    render(piece) {
        return piece;
    }
}
exports.TextChunk = TextChunk;
async function getTextContentBelowBudget(sizing, breakOn, fullText, cancellation) {
    if (breakOn instanceof RegExp) {
        if (!breakOn.global) {
            throw new Error(`\`breakOn\` expression must have the global flag set (got ${breakOn})`);
        }
        breakOn.lastIndex = 0;
    }
    let outputText = '';
    let lastIndex = -1;
    while (lastIndex < fullText.length) {
        let index;
        if (typeof breakOn === 'string') {
            index = fullText.indexOf(breakOn, lastIndex === -1 ? 0 : lastIndex + breakOn.length);
        }
        else {
            index = breakOn.exec(fullText)?.index ?? -1;
        }
        if (index === -1) {
            index = fullText.length;
        }
        const next = outputText + fullText.slice(Math.max(0, lastIndex), index);
        if ((await sizing.countTokens({ type: _1.Raw.ChatCompletionContentPartKind.Text, text: next }, cancellation)) > sizing.tokenBudget) {
            return outputText;
        }
        outputText = next;
        lastIndex = index;
    }
    return outputText;
}
class Image extends promptElement_1.PromptElement {
    constructor(props) {
        super(props);
    }
    render() {
        return vscpp(vscppf, null, this.props.children);
    }
}
exports.Image = Image;
/**
 * A document content part (e.g. PDF) that is a direct child of a {@link ChatMessagePromptElement}.
 */
class Document extends promptElement_1.PromptElement {
    constructor(props) {
        super(props);
    }
    render() {
        return vscpp(vscppf, null, this.props.children);
    }
}
exports.Document = Document;
/**
 * A utility for assigning priorities to a list of prompt elements.
 */
class PrioritizedList extends promptElement_1.PromptElement {
    render() {
        const { children, priority = 0, descending } = this.props;
        if (!children) {
            return;
        }
        return (vscpp(vscppf, null, children.map((child, i) => {
            if (!child) {
                return;
            }
            const thisPriority = descending
                ? // First element in array of children has highest priority
                    priority - i
                : // Last element in array of children has highest priority
                    priority - children.length + i;
            if (typeof child !== 'object') {
                return vscpp(TextChunk, { priority: thisPriority }, child);
            }
            child.props ??= {};
            child.props.priority = thisPriority;
            return child;
        })));
    }
}
exports.PrioritizedList = PrioritizedList;
/**
 * A utility to include the result of a tool called using the `vscode.lm.invokeTool` API.
 */
class ToolResult extends promptElement_1.PromptElement {
    render() {
        // note: future updates to content types should be handled here for backwards compatibility
        return (vscpp(vscppf, null, this.props.data.content.map(part => {
            if (part && typeof part.value === 'string') {
                return part.value;
            }
            else if (part &&
                part.value &&
                typeof part.value.node === 'object') {
                return (vscpp("elementJSON", { data: part.value }));
            }
        })));
    }
}
exports.ToolResult = ToolResult;
/**
 * Marker element that uses the legacy global prioritization algorithm (0.2.x
 * if this library) for pruning child elements. This will be removed in
 * the future.
 *
 * @deprecated
 */
class LegacyPrioritization extends promptElement_1.PromptElement {
    render() {
        return vscpp(vscppf, null, this.props.children);
    }
}
exports.LegacyPrioritization = LegacyPrioritization;
/**
 * Marker element that ensures all of its children are either included, or
 * not included. This is similar to the `<TextChunk />` element, but it is more
 * basic and can contain extrinsic children.
 */
class Chunk extends promptElement_1.PromptElement {
    render() {
        return vscpp(vscppf, null, this.props.children);
    }
}
exports.Chunk = Chunk;
/**
 * An element that can expand to fill the remaining token budget. Takes
 * a `value` function that is initially called with the element's token budget,
 * and may be called multiple times with the new token budget as the prompt
 * is resized.
 */
class Expandable extends promptElement_1.PromptElement {
    async render(_state, sizing) {
        return vscpp(vscppf, null, await this.props.value(sizing));
    }
}
exports.Expandable = Expandable;
/**
 * An element that ensures its children don't exceed a certain number of
 * `maxTokens`. Its contents are pruned to fit within the budget before
 * the overall prompt pruning is run.
 */
class TokenLimit extends promptElement_1.PromptElement {
    render() {
        return vscpp(vscppf, null, this.props.children);
    }
}
exports.TokenLimit = TokenLimit;
class AbstractKeepWith extends promptElement_1.PromptElement {
}
exports.AbstractKeepWith = AbstractKeepWith;
let keepWidthId = 0;
/**
 * Returns a PromptElement that ensures each wrapped element is retained only
 * so long as each other wrapped is not empty.
 *
 * This is useful when dealing with tool calls, for example. In that case,
 * your tool call request should only be rendered if the tool call response
 * survived prioritization. In that case, you implement a `render` function
 * like so:
 *
 * ```
 * render() {
 *   const KeepWith = useKeepWith();
 *   return <>
 *     <KeepWith priority={2}><ToolCallRequest>...</ToolCallRequest></KeepWith>
 *     <KeepWith priority={1}><ToolCallResponse>...</ToolCallResponse></KeepWith>
 *   </>;
 * }
 * ```
 *
 * Unlike `<Chunk />`, which blocks pruning of any child elements and simply
 * removes them as a block, `<KeepWith />` in this case will allow the
 * `ToolCallResponse` to be pruned, and if it's fully pruned it will also
 * remove the `ToolCallRequest`.
 */
function useKeepWith() {
    const id = keepWidthId++;
    return class KeepWith extends AbstractKeepWith {
        static id = id;
        id = id;
        render() {
            return vscpp(vscppf, null, this.props.children);
        }
    };
}
/**
 * An element that returns its `alt` prop if its children are empty at the
 * time when it's rendered. This is especially useful when you require
 * fallback logic for opaque child data, such as tool calls.
 */
class IfEmpty extends promptElement_1.PromptElement {
    render() {
        return (vscpp(vscppf, null,
            vscpp(LogicalWrapper, null, this.props.alt),
            vscpp(LogicalWrapper, { flexGrow: 1 }, this.props.children)));
    }
}
exports.IfEmpty = IfEmpty;
class LogicalWrapper extends promptElement_1.PromptElement {
    render() {
        return vscpp(vscppf, null, this.props.children);
    }
}
exports.LogicalWrapper = LogicalWrapper;
