AWS Comprehend Custom Classification in MCP Servers — Multi-class, Multi-label, Endpoints

Multi-class vs multi-label: choosing the right mode

The mode choice is architectural — it determines how Comprehend trains the model and how you interpret output scores:

For AliveMCP's use case — classifying user-submitted issue reports into tool handler categories — multi-class is the right choice: each issue should route to exactly one primary handler (server-unreachable, cert-expired, schema-drift, auth-failure, etc.). Multi-label would be appropriate for tagging a document with multiple applicable topics, like a blog post that covers both "monitoring" and "MCP protocol".

Training data: multi-class CSV format

The CSV training file has no header row. Each line contains exactly two comma-separated fields: the class label, then the document text. If the text contains commas, wrap it in double quotes:

SERVER_DOWN,"The mcp-server-filesystem endpoint is not responding. AliveMCP shows it has been down for 45 minutes."
SERVER_DOWN,"Getting connection refused errors on all calls to mcp-server-brave-search since 14:30 UTC."
AUTH_FAILURE,"mcp-server-github returns 401 on every tool call. Checked credentials and they look correct."
AUTH_FAILURE,"Authentication failing for mcp-server-gitlab. Token was rotated yesterday but the error persists."
SCHEMA_DRIFT,"mcp-server-postgres tool schema changed and now our orchestrator can't find the query_table tool."
SCHEMA_DRIFT,"The read_file tool on mcp-server-filesystem changed its parameter names in the last deployment."
CERT_EXPIRED,"SSL certificate error when connecting to mcp-server-everything. Browser also shows cert warning."
CERT_EXPIRED,"TLS handshake failing for mcp-server-memory, looks like the certificate expired two days ago."

Requirements for multi-class training data:

• Minimum 10 examples per class (100+ per class produces significantly better models)
• Class labels must not contain commas or newlines; spaces are allowed but lowercase is conventional
• Text must be plain UTF-8 — HTML entities, Unicode escapes, and control characters degrade model quality
• Each class should appear in a roughly balanced distribution; a 10:1 class imbalance significantly impacts minority class recall
• Maximum 1 million examples per training file; split into multiple files if larger

Training data: multi-label CSV format

Multi-label format is identical except the label field contains pipe-separated labels and should be quoted when multiple labels are present:

# Single label (no quotes required)
monitoring,"How do I set up uptime monitoring for my MCP server?"
# Multiple labels (pipe-separated, quoted)
"monitoring|alerting","Configure webhook alerts when an MCP server goes down."
"monitoring|metrics","What metrics does AliveMCP track for each endpoint?"
"alerting|slack","How do I connect AliveMCP alerts to a Slack channel?"

In multi-label mode, a document with no applicable labels is valid — you can include examples where the label field is empty (an empty string) to represent documents that should receive no tags. This helps the model learn to abstain rather than forcing a label onto every input.

Creating the document classifier

import {
  ComprehendClient,
  CreateDocumentClassifierCommand,
  DocumentClassifierMode,
  LanguageCode,
} from '@aws-sdk/client-comprehend';

const client = new ComprehendClient({ region: 'us-east-1' });

const { DocumentClassifierArn } = await client.send(
  new CreateDocumentClassifierCommand({
    DocumentClassifierName: 'mcp-issue-classifier-v1',
    LanguageCode: LanguageCode.EN,
    Mode: DocumentClassifierMode.MULTI_CLASS,
    DataAccessRoleArn: 'arn:aws:iam::123456789012:role/ComprehendDataAccessRole',
    InputDataConfig: {
      DataFormat: 'COMPREHEND_CSV',
      S3Uri: 's3://my-bucket/training/classifier-train.csv',
      // Optional: separate test file for evaluation metrics
      // TestS3Uri: 's3://my-bucket/training/classifier-test.csv',
      // Optional: augmented manifests from SageMaker Ground Truth
      // DataFormat: 'AUGMENTED_MANIFEST',
      // AugmentedManifests: [{ S3Uri: '...', AttributeNames: ['classification-label'] }],
    },
    OutputDataConfig: {
      // Evaluation metrics (confusion matrix, F1 per class) written here
      S3Uri: 's3://my-bucket/classifier-output/',
    },
    Tags: [{ Key: 'Project', Value: 'AliveMCP' }],
    // Version name for tracking iterations
    VersionName: 'v1-2026-11-03',
  })
);

The VersionName field is a freeform string that becomes part of the classifier ARN. Use semantic versioning or date strings — the ARN format is arn:aws:comprehend:region:account:document-classifier/name/version/versionName. Versioning lets you compare model quality across iterations without renaming the classifier: ListDocumentClassifierSummaries returns all versions for a given name.

Training takes 30–120 minutes depending on dataset size. Poll DescribeDocumentClassifier until Status is TRAINED. Check ClassifierMetadata.EvaluationMetrics for precision, recall, and F1 per class — these are computed on your test split (explicit test file or automatic 10% holdout).

Deploying a real-time endpoint

import { CreateEndpointCommand, DescribeEndpointCommand } from '@aws-sdk/client-comprehend';

const { EndpointArn } = await client.send(new CreateEndpointCommand({
  EndpointName: 'mcp-classifier-prod',
  ModelArn: documentClassifierArn,  // the trained classifier ARN
  DesiredInferenceUnits: 1,  // 1 unit = 100 chars/sec
  Tags: [{ Key: 'Project', Value: 'AliveMCP' }],
}));

// Wait for IN_SERVICE (5–15 minutes)
async function waitForEndpoint(arn: string): Promise<void> {
  while (true) {
    const { EndpointProperties: ep } = await client.send(
      new DescribeEndpointCommand({ EndpointArn: arn })
    );
    if (ep?.Status === 'IN_SERVICE') return;
    if (ep?.Status === 'FAILED') throw new Error(`Endpoint failed: ${ep.Message ?? ''}`);
    await new Promise(r => setTimeout(r, 30_000));
  }
}

Endpoint scaling: each inference unit handles 100 characters/second. For a typical MCP tool input of 500 characters, 1 inference unit handles ~12 concurrent classifications per second — sufficient for most single-server workloads. If classification latency under load exceeds 200ms, increase DesiredInferenceUnits via UpdateEndpointCommand (takes effect in 5–10 minutes).

Real-time classification in an MCP tool handler

import {
  ClassifyDocumentCommand,
  type DocumentClass,
  type DocumentLabel,
} from '@aws-sdk/client-comprehend';
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { z } from 'zod';

const server = new McpServer({ name: 'issue-router', version: '1.0.0' });
const ENDPOINT_ARN = process.env.COMPREHEND_CLASSIFIER_ARN!;

server.tool('classify_issue', {
  issue_text: z.string().max(4000),
}, async ({ issue_text }) => {
  const { Classes, Labels } = await client.send(new ClassifyDocumentCommand({
    Text: issue_text,
    EndpointArn: ENDPOINT_ARN,
  }));

  // Multi-class: Classes[] sorted by Score descending
  // Classes[0] is always the top prediction
  const topClass: DocumentClass | undefined = Classes?.[0];

  if (!topClass || (topClass.Score ?? 0) < 0.70) {
    // Low-confidence prediction — route to fallback handler
    return {
      content: [{
        type: 'text',
        text: JSON.stringify({
          classification: 'UNKNOWN',
          confidence: topClass?.Score ?? 0,
          message: 'Issue does not match known categories with sufficient confidence',
        }),
      }],
    };
  }

  // Multi-label example (if using MULTI_LABEL mode):
  // Labels[] contains all labels — filter by threshold
  // const tags = (Labels ?? []).filter(l => (l.Score ?? 0) >= 0.60).map(l => l.Name);

  return {
    content: [{
      type: 'text',
      text: JSON.stringify({
        classification: topClass.Name,
        confidence: topClass.Score,
        // Include top 3 classes for transparency
        top_classes: (Classes ?? []).slice(0, 3).map(c => ({
          name: c.Name,
          score: c.Score,
        })),
      }),
    }],
  };
});

For multi-class mode, Classes[] is pre-sorted by score descending — Classes[0] is always the top prediction. For multi-label mode, Labels[] is returned instead of Classes[], and the scores are independent (they do not sum to 1.0). Apply a per-label threshold (typically 0.50–0.70) to decide which labels to assign.

Async batch classification without an endpoint

import {
  StartDocumentClassificationJobCommand,
  DescribeDocumentClassificationJobCommand,
} from '@aws-sdk/client-comprehend';

// No endpoint required — cheaper for large batches
const { JobId } = await client.send(new StartDocumentClassificationJobCommand({
  JobName: `classify-batch-${Date.now()}`,
  DocumentClassifierArn: documentClassifierArn,
  DataAccessRoleArn: 'arn:aws:iam::123456789012:role/ComprehendDataAccessRole',
  InputDataConfig: {
    S3Uri: 's3://my-bucket/batch-input/',
    InputFormat: 'ONE_DOC_PER_LINE',
  },
  OutputDataConfig: {
    S3Uri: 's3://my-bucket/batch-output/',
  },
}));

// Output format: one JSON line per document
// {"Line": 0, "Classes": [{"Name": "SERVER_DOWN", "Score": 0.94}, ...]}  (MULTI_CLASS)
// {"Line": 0, "Labels": [{"Name": "monitoring", "Score": 0.88}, ...]}    (MULTI_LABEL)

Batch classification cost: $0.0005 per 100 characters processed (same as real-time, no endpoint surcharge). For periodic bulk classification tasks — like nightly triage of accumulated support tickets — batch jobs are significantly cheaper than keeping an endpoint running 24/7. A 24/7 endpoint at 1 inference unit costs ~$0.36/day in endpoint charges alone, independent of volume.

Model versioning and A/B testing classifiers

Each CreateDocumentClassifierCommand with the same DocumentClassifierName but different VersionName creates a new version under the same classifier name. List versions with ListDocumentClassifierSummaries. To A/B test a new classifier version:

1. Train the new version (VersionName: 'v2-2026-12-01')
2. Deploy it to a second endpoint (EndpointName: 'mcp-classifier-v2-canary')
3. Route a percentage of traffic to the canary endpoint in your MCP tool handler
4. Compare Classes[0].Score distributions and downstream routing accuracy between versions
5. Update the production endpoint to the new model ARN via UpdateEndpointCommand (in-place, no downtime)
6. Delete the canary endpoint

UpdateEndpointCommand accepts a new DesiredModelArn pointing at the new classifier version. The update takes 5–10 minutes; during the transition the endpoint continues serving traffic from the old model. No cold-start window or traffic interruption.