"use strict";
/**
 * Inserts aggregated trace DAG data into ClickHouse tables:
 * aggregated_trace_dag_metadata, aggregated_trace_dag_nodes, aggregated_trace_dag_edges,
 * and trusys_dev.trace_dag_metadata (per-trace metadata with tools, agents, models, metrics_state, plugins_state).
 * Schema includes metrics_state and plugins_state Map columns.
 */
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.buildDagWithPromptAndResponse = buildDagWithPromptAndResponse;
exports.getToolsAndAgentsFromSpans = getToolsAndAgentsFromSpans;
exports.insertAggregatedTraceDagData = insertAggregatedTraceDagData;
const dayjs_1 = __importDefault(require("dayjs"));
const frameworkParser_1 = require("./observe/frameworkParser");
;
;
const EMPTY_MAP_STATE_TYPE = "Map(String, Tuple(AggregateFunction(sum, UInt64), AggregateFunction(sum, UInt64), AggregateFunction(sum, Float64), AggregateFunction(sum, UInt64)))";
function chString(s) {
    if (s === null || s === undefined)
        return 'NULL';
    return `'${String(s).replace(/'/g, "''")}'`;
}
/** Format string array for ClickHouse Array(String) literal: ['a','b'] or []. */
function chArrayString(arr) {
    if (!Array.isArray(arr) || arr.length === 0)
        return '[]';
    return `[${arr.map((s) => chString(String(s ?? '').trim())).join(', ')}]`;
}
/** Get first non-empty session.id from span attributes. */
function getSessionIdFromSpans(spans) {
    if (!Array.isArray(spans))
        return '';
    for (const s of spans) {
        const id = s?.attributes?.['session.id'] ?? s?.attributes?.['session_id'];
        if (id != null && String(id).trim())
            return String(id).trim();
    }
    return '';
}
function buildMetricsMapStateExpr(entries) {
    const rows = Object.entries(entries);
    if (rows.length === 0) {
        return `CAST(map(), '${EMPTY_MAP_STATE_TYPE}')`;
    }
    const parts = [];
    for (const [k, v] of rows) {
        parts.push(chString(k));
        parts.push(`(` +
            `sumState(toUInt64(${Math.max(0, Math.floor(v.pass))})), ` +
            `sumState(toUInt64(${Math.max(0, Math.floor(v.fail))})), ` +
            `sumState(toFloat64(${Number.isFinite(v.scoreSum) ? v.scoreSum : 0})), ` +
            `sumState(toUInt64(${Math.max(0, Math.floor(v.scoreCount))}))` +
            `)`);
    }
    return `map(${parts.join(', ')})`;
}
/** Extract plugin id from test_case: metadata.pluginId, or from assert type "trusys:redteam:contracts" → "contracts", "trusys:redteam:bias:age" → "bias:age". */
function getPluginIdFromTestCase(testCase, row) {
    const tc = testCase ?? row?.testCase ?? row?.test_case;
    const fromMeta = tc?.metadata?.pluginId ?? tc?.metadata?.plugin_id ?? row?.metadata?.pluginId ?? row?.metadata?.plugin_id;
    if (fromMeta && String(fromMeta).trim())
        return String(fromMeta).trim();
    const assertType = tc?.assert?.[0]?.type ?? tc?.assert?.[0]?.metric;
    if (assertType && String(assertType).startsWith('trusys:redteam:')) {
        const suffix = String(assertType).slice('trusys:redteam:'.length).trim();
        return suffix || 'security';
    }
    return 'security';
}
/** True if this row is a plugin (redteam) eval: test_case has metadata.pluginId or assert type trusys:redteam:* or is_redteam=1. */
function isPluginRow(row) {
    if (Number(row?.is_redteam ?? row?.isRedteam ?? 0) === 1)
        return true;
    const tc = row?.testCase ?? row?.test_case;
    if (tc?.metadata?.pluginId ?? tc?.metadata?.plugin_id)
        return true;
    const assertType = tc?.assert?.[0]?.type ?? tc?.assert?.[0]?.metric;
    return assertType != null && String(assertType).startsWith('trusys:redteam');
}
function buildMetricsAndPluginsMapsFromEvalResults(dbResults) {
    const metricsMap = {};
    const pluginsMap = {};
    if (!Array.isArray(dbResults) || dbResults.length === 0) {
        return { metricsMap, pluginsMap };
    }
    for (const r of dbResults ?? []) {
        const testCase = r?.testCase ?? r?.test_case;
        const grading = r?.grading_result ?? r?.gradingResult;
        const isEmptyGrading = grading == null || (typeof grading === 'object' && Object.keys(grading).length === 0);
        if (isPluginRow(r)) {
            const pluginId = getPluginIdFromTestCase(testCase, r);
            const pass = !isEmptyGrading && grading?.pass === true ? 1 : 0;
            const fail = isEmptyGrading || grading?.pass !== true ? 1 : 0;
            const score = Number(grading?.score ?? r?.score ?? 0);
            const cur = pluginsMap[pluginId] ?? { pass: 0, fail: 0, scoreSum: 0, scoreCount: 0 };
            cur.pass += pass;
            cur.fail += fail;
            cur.scoreSum += Number.isFinite(score) ? score : 0;
            cur.scoreCount += 1;
            pluginsMap[pluginId] = cur;
            continue;
        }
        const asserts = testCase?.assert;
        if (!Array.isArray(asserts) || asserts.length === 0)
            continue;
        const componentResults = grading?.componentResults ?? [];
        const seenMetricsFromComponents = new Set();
        for (const comp of componentResults) {
            const assertion = comp?.assertion ?? comp;
            const metric = String(assertion?.type ?? assertion?.displayLabel ?? assertion?.metric ?? '').trim() || 'unknown';
            seenMetricsFromComponents.add(metric);
            const cur = metricsMap[metric] ?? { pass: 0, fail: 0, scoreSum: 0, scoreCount: 0 };
            cur.pass += comp?.pass === true ? 1 : 0;
            cur.fail += comp?.pass !== true ? 1 : 0;
            cur.scoreSum += Number.isFinite(Number(comp?.score)) ? Number(comp.score) : 0;
            cur.scoreCount += 1;
            metricsMap[metric] = cur;
        }
        for (const a of asserts) {
            const metric = String(a?.type ?? a?.metric ?? a?.name ?? '').trim() || 'unknown';
            if (seenMetricsFromComponents.has(metric))
                continue;
            const cur = metricsMap[metric] ?? { pass: 0, fail: 0, scoreSum: 0, scoreCount: 0 };
            cur.fail += 1;
            cur.scoreCount += 1;
            metricsMap[metric] = cur;
        }
    }
    return { metricsMap, pluginsMap };
}
/** Map parser node kind to insert node_type (prompt/agent/tool/model/response only; workflow dropped, other→agent). */
function nodeKindToInsertType(kind) {
    if (kind === 'prompt')
        return 'prompt';
    if (kind === 'response')
        return 'response';
    if (kind === 'agent')
        return 'agent';
    if (kind === 'tool')
        return 'tool';
    if (kind === 'model')
        return 'model';
    if (kind === 'workflow')
        return 'agent'; // should not appear in transformed nodes
    return 'agent'; // other → agent
}
function buildDagWithPromptAndResponse(spans, dag, promptEntityId, responseEntityId) {
    // Use the DAG data as-is from the parser - no extra processing needed
    const nodes = dag.nodes
        .filter((n) => n.kind !== 'workflow') // Only filter out workflow nodes
        .map((n) => ({
        ...(n ?? {}),
        id: n.id,
        kind: nodeKindToInsertType(n.kind),
        label: n.label ?? n.id,
    }));
    const edges = dag.edges.map((e) => ({
        ...(e ?? {}),
        from: e.from,
        to: e.to,
        type: e.type === 'invokes' ? 'invokes' : e.type === 'parent' ? 'parent' : 'sequence',
        count: e.count ?? 1,
    }));
    return { nodes, edges };
}
function edgeTypeToStr(type) {
    if (type === 'invokes')
        return 'invokes';
    if (type === 'sequence')
        return 'sequence';
    if (type === 'parent')
        return 'parent';
    return 'sequence';
}
/**
 * Build DAG from spans and return tool and agent node ids for this trace.
 * Used to populate eval_results.tools and eval_results.agents before insert.
 */
function getToolsAndAgentsFromSpans(spans) {
    const tools = [];
    const agents = [];
    if (!Array.isArray(spans) || spans.length === 0)
        return { tools, agents };
    const parser = new frameworkParser_1.MultiFrameworkParser();
    const { dag, promptEntityId, responseEntityId } = parser.getDagData(spans);
    if (!dag?.nodes?.length)
        return { tools, agents };
    const { nodes: transformedNodes } = buildDagWithPromptAndResponse(spans, dag, promptEntityId, responseEntityId);
    for (const n of transformedNodes) {
        if (n.kind === 'tool')
            tools.push(n.id?.replace('tool:', ''));
        else if (n.kind === 'agent')
            agents.push(n.id?.replace('agent:', ''));
    }
    return { tools, agents };
}
async function insertAggregatedTraceDagData(params, deps) {
    const { spans, traceId, projectId, workspaceId, applicationId, dbResults } = params;
    const { clickHouseService, logger } = deps;
    if (!Array.isArray(spans) || spans.length === 0 || !traceId) {
        return;
    }
    const parser = new frameworkParser_1.MultiFrameworkParser();
    const parsed = parser.parse(spans);
    const { dag, promptEntityId, responseEntityId } = parser.getDagData(spans);
    if (!dag?.nodes?.length) {
        return;
    }
    const environment = String(parsed?.environment ?? parsed?.metadata?.environment ?? '').trim();
    const version = String(parsed?.version ?? parsed?.metadata?.version ?? '').trim();
    const { nodes: transformedNodes, edges: transformedEdges } = buildDagWithPromptAndResponse(spans, dag, promptEntityId, responseEntityId);
    const timestamps = spans.map((s) => new Date(s.timestamp).getTime()).filter(Number.isFinite);
    const minTs = timestamps.length ? Math.min(...timestamps) : Date.now();
    const day = (0, dayjs_1.default)(minTs).format('YYYY-MM-DD');
    const { metricsMap, pluginsMap } = buildMetricsAndPluginsMapsFromEvalResults(dbResults);
    const metricsStateExpr = buildMetricsMapStateExpr(metricsMap);
    const pluginsStateExpr = buildMetricsMapStateExpr(pluginsMap);
    const totalNodes = transformedNodes.length;
    const totalEdges = transformedEdges.length;
    const metadataSelect = `SELECT
    toDate(${chString(day)}) as day,
    ${chString(projectId)} as project_id,
    ${chString(workspaceId)} as workspace_id,
    ${chString(applicationId)} as application_id,
    ${chString(environment)} as environment,
    ${chString(version)} as version,
    sumState(toUInt64(1)) as trace_count_state,
    maxState(toUInt32(${totalNodes})) as total_nodes_state,
    maxState(toUInt32(${totalEdges})) as total_edges_state,
    ${metricsStateExpr} as metrics_state,
    ${pluginsStateExpr} as plugins_state`;
    const metadataSql = `INSERT INTO aggregated_trace_dag_metadata (day, project_id, workspace_id, application_id, environment, version, trace_count_state, total_nodes_state, total_edges_state, metrics_state, plugins_state) ${metadataSelect}`;
    await clickHouseService.command(metadataSql);
    const nodeSelects = transformedNodes.map((node) => {
        return `SELECT
      toDate(${chString(day)}) as day,
      ${chString(projectId)} as project_id,
      ${chString(workspaceId)} as workspace_id,
      ${chString(applicationId)} as application_id,
      ${chString(environment)} as environment,
      ${chString(version)} as version,
      ${chString(node.id)} as node_name,
      ${chString(node.label)} as node_label,
      ${chString(node.kind)} as node_type,
      sumState(toUInt64(1)) as call_count_state,
      ${metricsStateExpr} as metrics_state,
      ${pluginsStateExpr} as plugins_state`;
    });
    const nodesSql = `INSERT INTO aggregated_trace_dag_nodes (day, project_id, workspace_id, application_id, environment, version, node_name, node_label, node_type, call_count_state, metrics_state, plugins_state) ${nodeSelects.join(' UNION ALL ')}`;
    await clickHouseService.command(nodesSql);
    if (transformedEdges.length > 0) {
        const edgeSelects = transformedEdges.map((edge) => {
            const edgeType = edgeTypeToStr(edge.type);
            return `SELECT
        toDate(${chString(day)}) as day,
        ${chString(projectId)} as project_id,
        ${chString(workspaceId)} as workspace_id,
        ${chString(applicationId)} as application_id,
        ${chString(environment)} as environment,
        ${chString(version)} as version,
        ${chString(edge.from)} as from_node,
        ${chString(edge.to)} as to_node,
        ${chString(edgeType)} as edge_type,
        sumState(toUInt64(${edge.count})) as call_count_state,
        ${metricsStateExpr} as metrics_state,
        ${pluginsStateExpr} as plugins_state`;
        });
        const edgesSql = `INSERT INTO aggregated_trace_dag_edges (day, project_id, workspace_id, application_id, environment, version, from_node, to_node, edge_type, call_count_state, metrics_state, plugins_state) ${edgeSelects.join(' UNION ALL ')}`;
        await clickHouseService.command(edgesSql);
    }
    // Insert per-trace row into trusys_dev.trace_dag_metadata (tools, agents, models; metrics_state and plugins_state only when dbResults present)
    const sessionId = getSessionIdFromSpans(spans);
    const { tools, agents } = getToolsAndAgentsFromSpans(spans);
    const modelsFromParser = parser.getModels(spans);
    const models = [
        ...new Set((modelsFromParser ?? []).map((m) => String(m?.modelName ?? m?.id ?? '').trim()).filter(Boolean)),
    ];
    const hasDbResults = Array.isArray(dbResults) && dbResults.length > 0;
    const traceDagMetricsState = hasDbResults ? metricsStateExpr : `CAST(map(), '${EMPTY_MAP_STATE_TYPE}')`;
    const traceDagPluginsState = hasDbResults ? pluginsStateExpr : `CAST(map(), '${EMPTY_MAP_STATE_TYPE}')`;
    const traceDagMetadataSelect = `SELECT
    ${chString(projectId)} as project_id,
    ${chString(workspaceId)} as workspace_id,
    ${chString(applicationId)} as application_id,
    ${chString(environment)} as environment,
    ${chString(version)} as version,
    ${chString(traceId)} as traceId,
    ${chString(sessionId || '')} as sessionId,
    ${chArrayString(tools)} as tools,
    ${chArrayString(agents)} as agents,
    ${chArrayString(models)} as models,
    ${traceDagMetricsState} as metrics_state,
    ${traceDagPluginsState} as plugins_state`;
    const traceDagMetadataSql = `INSERT INTO trusys_dev.trace_dag_metadata (project_id, workspace_id, application_id, environment, version, traceId, sessionId, tools, agents, models, metrics_state, plugins_state) ${traceDagMetadataSelect}`;
    await clickHouseService.command(traceDagMetadataSql);
}
//# sourceMappingURL=aggregatedTraceDagInsert.js.map