"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.ContextAttributes = void 0;
exports.setPromptTemplate = setPromptTemplate;
exports.clearPromptTemplate = clearPromptTemplate;
exports.getPromptTemplate = getPromptTemplate;
exports.setSession = setSession;
exports.clearSession = clearSession;
exports.getSession = getSession;
exports.setMetadata = setMetadata;
exports.clearMetadata = clearMetadata;
exports.getMetadata = getMetadata;
exports.setUser = setUser;
exports.clearUser = clearUser;
exports.getUser = getUser;
exports.setTags = setTags;
exports.clearTags = clearTags;
exports.getTags = getTags;
exports.setAttributes = setAttributes;
exports.clearAttributes = clearAttributes;
exports.getAttributes = getAttributes;
exports.getAttributesFromContext = getAttributesFromContext;
const openinference_semantic_conventions_1 = require("@arizeai/openinference-semantic-conventions");
const api_1 = require("@opentelemetry/api");
const utils_1 = require("../utils");
const core_1 = require("@opentelemetry/core");
const CONTEXT_ATTRIBUTES_ATTRIBUTES_KEY = "attributes";
exports.ContextAttributes = {
    [openinference_semantic_conventions_1.PROMPT_TEMPLATE_TEMPLATE]: (0, api_1.createContextKey)(`OpenInference SDK Context Key ${openinference_semantic_conventions_1.PROMPT_TEMPLATE_TEMPLATE}`),
    [openinference_semantic_conventions_1.PROMPT_TEMPLATE_VARIABLES]: (0, api_1.createContextKey)(`OpenInference SDK Context Key ${openinference_semantic_conventions_1.PROMPT_TEMPLATE_VARIABLES}`),
    [openinference_semantic_conventions_1.PROMPT_TEMPLATE_VERSION]: (0, api_1.createContextKey)(`OpenInference SDK Context Key ${openinference_semantic_conventions_1.PROMPT_TEMPLATE_VERSION}`),
    [openinference_semantic_conventions_1.SESSION_ID]: (0, api_1.createContextKey)(`OpenInference SDK Context Key ${openinference_semantic_conventions_1.SESSION_ID}`),
    [openinference_semantic_conventions_1.METADATA]: (0, api_1.createContextKey)(`OpenInference SDK Context Key ${openinference_semantic_conventions_1.METADATA}`),
    [openinference_semantic_conventions_1.USER_ID]: (0, api_1.createContextKey)(`OpenInference SDK Context Key ${openinference_semantic_conventions_1.USER_ID}`),
    [openinference_semantic_conventions_1.TAG_TAGS]: (0, api_1.createContextKey)(`OpenInference SDK Context Key ${openinference_semantic_conventions_1.TAG_TAGS}`),
    [CONTEXT_ATTRIBUTES_ATTRIBUTES_KEY]: (0, api_1.createContextKey)(`OpenInference SDK Context Key attributes`),
};
const { [openinference_semantic_conventions_1.PROMPT_TEMPLATE_TEMPLATE]: PROMPT_TEMPLATE_TEMPLATE_KEY, [openinference_semantic_conventions_1.PROMPT_TEMPLATE_VARIABLES]: PROMPT_TEMPLATE_VARIABLES_KEY, [openinference_semantic_conventions_1.PROMPT_TEMPLATE_VERSION]: PROMPT_TEMPLATE_VERSION_KEY, [openinference_semantic_conventions_1.SESSION_ID]: SESSION_ID_KEY, [openinference_semantic_conventions_1.METADATA]: METADATA_KEY, [openinference_semantic_conventions_1.USER_ID]: USER_ID_KEY, [openinference_semantic_conventions_1.TAG_TAGS]: TAG_TAGS_KEY, [CONTEXT_ATTRIBUTES_ATTRIBUTES_KEY]: ATTRIBUTES_KEY, } = exports.ContextAttributes;
function setPromptTemplate(context, attributes) {
    const { template, variables, version } = attributes;
    context = context.setValue(PROMPT_TEMPLATE_TEMPLATE_KEY, template);
    if (variables) {
        context = context.setValue(PROMPT_TEMPLATE_VARIABLES_KEY, (0, utils_1.safelyJSONStringify)(variables));
    }
    if (version) {
        context = context.setValue(PROMPT_TEMPLATE_VERSION_KEY, version);
    }
    return context;
}
function clearPromptTemplate(context) {
    context = context.deleteValue(PROMPT_TEMPLATE_TEMPLATE_KEY);
    context = context.deleteValue(PROMPT_TEMPLATE_VARIABLES_KEY);
    context = context.deleteValue(PROMPT_TEMPLATE_VERSION_KEY);
    return context;
}
function getPromptTemplate(context) {
    const maybeTemplate = context.getValue(PROMPT_TEMPLATE_TEMPLATE_KEY);
    const maybeVariables = context.getValue(PROMPT_TEMPLATE_VARIABLES_KEY);
    const maybeVersion = context.getValue(PROMPT_TEMPLATE_VERSION_KEY);
    const attributes = {};
    if (typeof maybeTemplate === "string") {
        attributes.template = maybeTemplate;
    }
    if (typeof maybeVariables === "string") {
        const parsedVariables = (0, utils_1.safelyJSONParse)(maybeVariables);
        attributes.variables = (0, utils_1.isObjectWithStringKeys)(parsedVariables)
            ? parsedVariables
            : undefined;
    }
    if (typeof maybeVersion === "string") {
        attributes.version = maybeVersion;
    }
    if (Object.keys(attributes).length === 0) {
        return;
    }
    return attributes;
}
function setSession(context, attributes) {
    const { sessionId } = attributes;
    return context.setValue(SESSION_ID_KEY, sessionId);
}
function clearSession(context) {
    return context.deleteValue(SESSION_ID_KEY);
}
/**
 * Retrieves the session ID from the given context.
 * @param context - The context object.
 * @returns {string | undefined} The session ID if it exists, otherwise undefined.
 */
function getSession(context) {
    const maybeSessionId = context.getValue(SESSION_ID_KEY);
    if (typeof maybeSessionId === "string") {
        return { sessionId: maybeSessionId };
    }
}
function setMetadata(context, attributes) {
    return context.setValue(METADATA_KEY, (0, utils_1.safelyJSONStringify)(attributes));
}
function clearMetadata(context) {
    return context.deleteValue(METADATA_KEY);
}
function getMetadata(context) {
    const maybeMetadata = context.getValue(METADATA_KEY);
    if (typeof maybeMetadata === "string") {
        const parsedMetadata = (0, utils_1.safelyJSONParse)(maybeMetadata);
        return (0, utils_1.isObjectWithStringKeys)(parsedMetadata) ? parsedMetadata : undefined;
    }
}
function setUser(context, attributes) {
    const { userId } = attributes;
    return context.setValue(USER_ID_KEY, userId);
}
function clearUser(context) {
    return context.deleteValue(USER_ID_KEY);
}
function getUser(context) {
    const maybeUserId = context.getValue(USER_ID_KEY);
    if (typeof maybeUserId === "string") {
        return { userId: maybeUserId };
    }
}
function setTags(context, attributes) {
    return context.setValue(TAG_TAGS_KEY, (0, utils_1.safelyJSONStringify)(attributes));
}
function clearTags(context) {
    return context.deleteValue(TAG_TAGS_KEY);
}
function getTags(context) {
    const maybeTags = context.getValue(TAG_TAGS_KEY);
    if (typeof maybeTags === "string") {
        const parsedTags = (0, utils_1.safelyJSONParse)(maybeTags);
        return (0, utils_1.isStringArray)(parsedTags) ? parsedTags : undefined;
    }
}
function setAttributes(context, attributes) {
    return context.setValue(ATTRIBUTES_KEY, (0, utils_1.safelyJSONStringify)(attributes));
}
function clearAttributes(context) {
    return context.deleteValue(ATTRIBUTES_KEY);
}
function getAttributes(context) {
    const maybeAttributes = context.getValue(ATTRIBUTES_KEY);
    if (typeof maybeAttributes === "string") {
        const parsedAttributes = (0, utils_1.safelyJSONParse)(maybeAttributes);
        return (0, utils_1.isAttributes)(parsedAttributes) ? parsedAttributes : undefined;
    }
}
/**
 * Gets the OpenInference attributes from the given context
 * @param context
 * @example span.setAttributes(getAttributesFromContext(context.active()));
 * @returns {Attributes} The OpenInference attributes formatted as OpenTelemetry span attributes.
 */
function getAttributesFromContext(context) {
    let attributes = {};
    Object.entries(exports.ContextAttributes).forEach(([key, symbol]) => {
        const maybeValue = context.getValue(symbol);
        if (key === CONTEXT_ATTRIBUTES_ATTRIBUTES_KEY) {
            if (typeof maybeValue === "string") {
                const parsedAttributes = (0, utils_1.safelyJSONParse)(maybeValue);
                if ((0, utils_1.isAttributes)(parsedAttributes)) {
                    attributes = Object.assign(Object.assign({}, attributes), parsedAttributes);
                }
            }
            return;
        }
        if ((0, core_1.isAttributeValue)(maybeValue) && maybeValue !== undefined) {
            attributes[key] = maybeValue;
        }
    });
    return attributes;
}
//# sourceMappingURL=contextAttributes.js.map