learning_ai_common_plat/services/mcp-server/src/modules/a2a/triage-quality-pipeline.ts

/**
 * TriageQualityAgent — A2A pipeline for MindLyst memory triage quality improvement.
 *
 * Agent roster (3 steps):
 *   1. LowConfidenceCollectorAgent — list memory items, filter for confidenceScore below threshold
 *   2. RetriageAgent               — re-run extraction on each low-confidence item
 *   3. TriageQualityReportAgent    — compare old vs new routing, auto-reassign if brainId changed
 *
 * MCP tools:
 *   mindlyst.memory.triageQuality(confidenceThreshold?, dryRun?) — run pipeline
 */

import { randomUUID } from 'node:crypto';
import { z } from 'zod';
import { registerTool } from '../tools/registry.js';
import type { McpToolRequest } from '../tools/types.js';
import {
  mindlystMemoryList,
  mindlystMemoryRetriage,
  mindlystMemoryReassign,
} from '../../lib/mindlyst-client.js';
import { config } from '../../lib/config.js';

// ── Types ──────────────────────────────────────────────────────────────────────

interface LowConfidenceItem {
  id: string;
  confidenceScore: number;
  currentBrainId: string | null;
  contentType: string;
}

interface RetriageResult {
  itemId: string;
  oldConfidenceScore: number;
  newConfidenceScore: number;
  oldBrainId: string | null;
  newBrainId: string | null;
  brainChanged: boolean;
  reassigned: boolean;
  error?: string;
}

export interface TriageQualityReport {
  runId: string;
  productId: 'mindlyst';
  confidenceThreshold: number;
  dryRun: boolean;
  totalFetched: number;
  lowConfidenceCount: number;
  retriaged: number;
  improved: number;
  brainChanges: number;
  reassigned: number;
  failed: number;
  perItem: RetriageResult[];
  summary: string;
  generatedAt: string;
}

// ── Step 1: LowConfidenceCollectorAgent ───────────────────────────────────────

async function collectLowConfidenceItems(
  confidenceThreshold: number,
  opts: { token?: string; requestId?: string }
): Promise<{ items: LowConfidenceItem[]; totalFetched: number }> {
  try {
    const result = await mindlystMemoryList({ limit: config.QUERY_MAX_LIMIT }, opts);

    const allItems = result.items;
    const lowConfidence = allItems
      .filter(item => item.triageResult.confidenceScore < confidenceThreshold)
      .map(item => ({
        id: item.id,
        confidenceScore: item.triageResult.confidenceScore,
        currentBrainId: item.brainIds[0] ?? null,
        contentType: item.triageResult.contentType,
      }));

    return { items: lowConfidence, totalFetched: allItems.length };
  } catch {
    return { items: [], totalFetched: 0 };
  }
}

// ── Step 2: RetriageAgent ─────────────────────────────────────────────────────

async function retriageItem(
  item: LowConfidenceItem,
  dryRun: boolean,
  opts: { token?: string; requestId?: string }
): Promise<RetriageResult> {
  if (dryRun) {
    return {
      itemId: item.id,
      oldConfidenceScore: item.confidenceScore,
      newConfidenceScore: item.confidenceScore,
      oldBrainId: item.currentBrainId,
      newBrainId: item.currentBrainId,
      brainChanged: false,
      reassigned: false,
    };
  }

  try {
    const updated = await mindlystMemoryRetriage(item.id, opts);
    const newConfidence = updated.triageResult.confidenceScore;
    const newBrainId = updated.triageResult.suggestedBrainId ?? updated.brainIds[0] ?? null;
    const brainChanged = newBrainId !== null && newBrainId !== item.currentBrainId;

    let reassigned = false;
    if (brainChanged && newBrainId) {
      try {
        await mindlystMemoryReassign(item.id, newBrainId, opts);
        reassigned = true;
      } catch {
        // best-effort reassignment
      }
    }

    return {
      itemId: item.id,
      oldConfidenceScore: item.confidenceScore,
      newConfidenceScore: newConfidence,
      oldBrainId: item.currentBrainId,
      newBrainId,
      brainChanged,
      reassigned,
    };
  } catch (err) {
    return {
      itemId: item.id,
      oldConfidenceScore: item.confidenceScore,
      newConfidenceScore: item.confidenceScore,
      oldBrainId: item.currentBrainId,
      newBrainId: item.currentBrainId,
      brainChanged: false,
      reassigned: false,
      error: err instanceof Error ? err.message : String(err),
    };
  }
}

// ── Step 3: TriageQualityReportAgent ──────────────────────────────────────────

function buildTriageReport(
  runId: string,
  confidenceThreshold: number,
  dryRun: boolean,
  totalFetched: number,
  results: RetriageResult[]
): TriageQualityReport {
  const now = new Date().toISOString();
  const improved = results.filter(
    r => r.newConfidenceScore > r.oldConfidenceScore && !r.error
  ).length;
  const brainChanges = results.filter(r => r.brainChanged).length;
  const reassigned = results.filter(r => r.reassigned).length;
  const failed = results.filter(r => !!r.error).length;

  const summary = dryRun
    ? `DRY RUN: ${results.length} items below confidence threshold ${confidenceThreshold} found. No retriaging performed.`
    : results.length === 0
      ? `All ${totalFetched} memory items meet confidence threshold ${confidenceThreshold}. No action needed.`
      : `Retriaged ${results.length - failed}/${results.length} items. ${improved} improved confidence, ${brainChanges} brain routing changes, ${reassigned} items reassigned. ${failed} failed.`;

  return {
    runId,
    productId: 'mindlyst',
    confidenceThreshold,
    dryRun,
    totalFetched,
    lowConfidenceCount: results.length,
    retriaged: results.length - failed,
    improved,
    brainChanges,
    reassigned,
    failed,
    perItem: results,
    summary,
    generatedAt: now,
  };
}

// ── Pipeline runner ────────────────────────────────────────────────────────────

async function runTriageQualityPipeline(
  confidenceThreshold: number,
  dryRun: boolean,
  req: McpToolRequest
): Promise<TriageQualityReport> {
  const runId = randomUUID();
  const opts = { token: req.headers.authorization?.slice(7), requestId: req.id };

  req.log.info(
    { runId, stepId: 'collect', confidenceThreshold, dryRun },
    'LowConfidenceCollectorAgent start'
  );
  const { items, totalFetched } = await collectLowConfidenceItems(confidenceThreshold, opts);
  req.log.info(
    { runId, stepId: 'collect', totalFetched, lowConfidenceCount: items.length },
    'LowConfidenceCollectorAgent done'
  );

  req.log.info({ runId, stepId: 'retriage', count: items.length, dryRun }, 'RetriageAgent start');
  const results: RetriageResult[] = [];
  for (const item of items) {
    const result = await retriageItem(item, dryRun, opts);
    results.push(result);
  }
  req.log.info(
    {
      runId,
      stepId: 'retriage',
      improved: results.filter(r => r.newConfidenceScore > r.oldConfidenceScore).length,
    },
    'RetriageAgent done'
  );

  req.log.info({ runId, stepId: 'report' }, 'TriageQualityReportAgent start');
  const report = buildTriageReport(runId, confidenceThreshold, dryRun, totalFetched, results);
  req.log.info(
    { runId, stepId: 'report', summary: report.summary },
    'TriageQualityReportAgent done'
  );

  return report;
}

// ── MCP tool registration ─────────────────────────────────────────────────────

registerTool({
  name: 'mindlyst.memory.triageQuality',
  description:
    'A2A pipeline: finds MindLyst memory items with confidenceScore below a threshold, re-runs extraction on each, and auto-reassigns items where brain routing has changed. Returns a quality improvement report with confidence deltas and brain reassignment counts. Use dryRun=true to preview. Requires admin role.',
  requiredRole: 'admin',
  inputSchema: z.object({
    confidenceThreshold: z.coerce
      .number()
      .min(0)
      .max(1)
      .default(0.5)
      .describe('Items with confidenceScore below this are re-triaged (default 0.5)'),
    dryRun: z
      .boolean()
      .default(false)
      .describe('If true, collect and count only — do not retriage or reassign'),
  }),
  async execute(args, req) {
    return runTriageQualityPipeline(args.confidenceThreshold, args.dryRun, req);
  },
});