AWS Comprehend Custom Entity Recognition in MCP Servers — Training, Inference, Endpoints

When custom entity recognition is the right tool

Comprehend's built-in DetectEntities API identifies general named entities across standard categories. It will correctly identify "AWS" as an ORGANIZATION and "New York" as a LOCATION, but it cannot identify "alivemcp-status-server" as an MCP_SERVER_SLUG or "CVE-2026-12345" as a VULNERABILITY_ID — those are domain concepts it has never seen.

Custom entity recognizers fill this gap. You train a model on documents you have annotated with your domain types, and the resulting recognizer learns to identify those patterns in new text. Comprehend supports up to 25 custom entity types per recognizer. Designing this list well matters: too broad (e.g., a single PRODUCT type for all product-related text) and precision suffers; too narrow (separate types for every product variant) and you need far more training examples per type. A workable rule: if you would write separate code to handle two entity types differently, they warrant separate types.

For MCP server teams, custom entity recognition is useful for:

• Extracting structured fields from freeform tool call arguments (e.g., identifying server slugs, version strings, or registry names in natural-language queries)
• Parsing domain-specific documents that MCP tools retrieve and return to AI orchestrators (contract clauses, log entries, technical specifications)
• Pre-processing text before storing it to enable structured querying downstream

Training data format: CSV annotations

The CSV annotation format requires two S3 objects: a documents file (one document per line, plain text) and an annotations file (tab-separated, one annotation span per row). The annotations file must have exactly this header row:

File	Line	Begin Offset	End Offset	Type

File is the S3 key of the documents file. Line is the 0-based line number in that file (each line is one document). Begin Offset and End Offset are character offsets within that line (0-based, end offset is exclusive). Type must exactly match one of the entity type names declared in your EntityTypes list.

# documents.txt (one document per line, stored in S3)
The mcp-server-filesystem endpoint has been down for 3 hours according to the AliveMCP dashboard.
User reported that mcp-server-brave-search returned TOOL_ERROR on all calls since the deployment.
AliveMCP shows mcp-server-sequential-thinking at 99.9% uptime over the last 30 days.

# annotations.csv (tab-separated)
File	Line	Begin Offset	End Offset	Type
s3://my-bucket/train/documents.txt	0	4	26	MCP_SERVER_SLUG
s3://my-bucket/train/documents.txt	1	18	41	MCP_SERVER_SLUG
s3://my-bucket/train/documents.txt	2	9	38	MCP_SERVER_SLUG

Annotation quality requirements: at least 200 annotated documents per entity type (more is better — 1,000+ per type produces significantly higher F1 scores). The training set must contain at least one positive example of every declared entity type; a type with zero annotations causes the training job to fail. Documents can be multi-sentence; the annotation offsets must precisely cover the entity span — off-by-one errors in offsets are the most common cause of poor model quality.

Training data format: augmented manifests

Augmented manifests are the output format of SageMaker Ground Truth labeling jobs. Each line is a JSON object containing the document text and the annotation label from the labeling job attribute name:

{"source": "The mcp-server-filesystem endpoint has been down for 3 hours.", "entity-labels": {"annotations": {"entities": [{"startOffset": 4, "endOffset": 26, "label": "MCP_SERVER_SLUG"}], "labels": [{"label": "MCP_SERVER_SLUG"}]}}}
{"source": "mcp-server-brave-search returned TOOL_ERROR on all calls.", "entity-labels": {"annotations": {"entities": [{"startOffset": 0, "endOffset": 23, "label": "MCP_SERVER_SLUG"}], "labels": [{"label": "MCP_SERVER_SLUG"}]}}}

When using augmented manifests, specify DataFormat: "AUGMENTED_MANIFEST" in the InputDataConfig and set AttributeNames to the labeling job attribute name (here: ["entity-labels"]). Augmented manifests support train/test split via separate manifest files for Documents (train) and TestDocuments (evaluation) — use at minimum an 80/20 split. Without an explicit test set, Comprehend performs automatic evaluation on a random 10% split of your training data.

Creating the entity recognizer

import {
  ComprehendClient,
  CreateEntityRecognizerCommand,
  LanguageCode,
  type CreateEntityRecognizerRequest,
} from '@aws-sdk/client-comprehend';

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

const params: CreateEntityRecognizerRequest = {
  RecognizerName: 'mcp-server-entities-v1',
  LanguageCode: LanguageCode.EN,
  DataAccessRoleArn: 'arn:aws:iam::123456789012:role/ComprehendDataAccessRole',
  InputDataConfig: {
    DataFormat: 'COMPREHEND_CSV',
    EntityTypes: [
      { Type: 'MCP_SERVER_SLUG' },
      { Type: 'REGISTRY_NAME' },
      { Type: 'TOOL_ERROR_CODE' },
    ],
    Documents: {
      S3Uri: 's3://my-bucket/train/documents.txt',
    },
    Annotations: {
      S3Uri: 's3://my-bucket/train/annotations.csv',
    },
    // Optional: separate evaluation set
    // TestDocuments: { S3Uri: 's3://my-bucket/test/documents.txt' },
    // TestAnnotations: { S3Uri: 's3://my-bucket/test/annotations.csv' },
  },
  Tags: [{ Key: 'Project', Value: 'AliveMCP' }],
};

const { EntityRecognizerArn } = await client.send(
  new CreateEntityRecognizerCommand(params)
);

The DataAccessRoleArn must have a trust policy allowing comprehend.amazonaws.com to assume it, with an aws:SourceAccount condition to prevent confused deputy attacks. The role needs s3:GetObject and s3:ListBucket on the training data bucket, and s3:PutObject if you specify an output location for evaluation results.

Entity type names must match the regex [A-Z_]+ — uppercase letters and underscores only, no spaces, digits, or hyphens. The recognizer ARN returned immediately; the training job runs asynchronously. Poll DescribeEntityRecognizer until Status is TRAINED (typically 30–90 minutes for training sets under 100,000 documents).

Monitoring training and evaluating model quality

import {
  DescribeEntityRecognizerCommand,
  type EntityRecognizerEvaluationMetrics,
} from '@aws-sdk/client-comprehend';

async function waitForTraining(arn: string): Promise<EntityRecognizerEvaluationMetrics> {
  while (true) {
    const { EntityRecognizerProperties: props } = await client.send(
      new DescribeEntityRecognizerCommand({ EntityRecognizerArn: arn })
    );

    const status = props?.Status;
    if (status === 'TRAINED') {
      return props.RecognizerMetadata?.EvaluationMetrics ?? {};
    }
    if (status === 'IN_ERROR') {
      throw new Error(`Training failed: ${props?.Message ?? 'unknown'}`);
    }

    // Poll every 60 seconds during training
    await new Promise(r => setTimeout(r, 60_000));
  }
}

// Evaluation metrics returned after TRAINED status:
// {
//   Precision: 0.87,      // of all predicted spans, what fraction were correct
//   Recall: 0.82,         // of all true spans, what fraction did we find
//   F1Score: 0.845,       // harmonic mean — primary metric to optimize
//   // Per-entity-type breakdown available in RecognizerMetadata.EntityTypes[]
// }

F1 score ≥ 0.80 is a reasonable production threshold for most use cases. If F1 is below 0.70, the most effective remedies in order of impact: (1) increase annotation volume for under-performing entity types — check per-type metrics in RecognizerMetadata.EntityTypes[]; (2) review annotation consistency — inter-annotator disagreement above 10% degrades model quality significantly; (3) check annotation offset precision — a systematic off-by-one error (missing leading/trailing character) is a common root cause; (4) add negative examples (documents that contain similar patterns but are NOT the target entity type).

Deploying a real-time endpoint

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

// Deploy the trained recognizer to a real-time endpoint
// 1 inference unit = 100 characters/second throughput, billed at ~$0.0005/hour
const { EndpointArn } = await client.send(new CreateEndpointCommand({
  EndpointName: 'mcp-entities-prod',
  ModelArn: entityRecognizerArn,
  DesiredInferenceUnits: 1,
  Tags: [{ Key: 'Project', Value: 'AliveMCP' }],
}));

// Wait for endpoint to become 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));
  }
}

Inference unit pricing applies from the moment the endpoint enters IN_SERVICE until it is deleted — there is no pause/stop. For development and low-throughput production workloads, prefer async batch jobs (no endpoint required) over real-time endpoints. Deploy real-time endpoints only when your MCP server must return entity extraction results within the same tool call (synchronous use case) and the document volume justifies the continuous cost.

Real-time inference from an MCP tool handler

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

const server = new McpServer({ name: 'entity-extractor', version: '1.0.0' });

const ENDPOINT_ARN = process.env.COMPREHEND_ENDPOINT_ARN!;
const CONFIDENCE_THRESHOLD = 0.80;

server.tool('extract_mcp_entities', {
  text: z.string().max(5000),  // DetectEntities limit: 5,000 bytes UTF-8
}, async ({ text }) => {
  // Comprehend measures the 5,000-byte limit in UTF-8 bytes, not characters
  const encoded = new TextEncoder().encode(text);
  if (encoded.length > 5000) {
    return { content: [{ type: 'text', text: JSON.stringify({
      error: 'Text exceeds 5,000 UTF-8 bytes — use batch job for longer documents'
    })}] };
  }

  const { Entities } = await client.send(new DetectEntitiesCommand({
    Text: text,
    EndpointArn: ENDPOINT_ARN,
    // LanguageCode is inferred from the recognizer — omit for custom recognizers
  }));

  const filtered: Entity[] = (Entities ?? []).filter(
    e => (e.Score ?? 0) >= CONFIDENCE_THRESHOLD
  );

  const result = filtered.map(e => ({
    text: e.Text,
    type: e.Type,
    score: e.Score,
    beginOffset: e.BeginOffset,
    endOffset: e.EndOffset,
  }));

  return { content: [{ type: 'text', text: JSON.stringify(result) }] };
});

The 5,000-byte limit is enforced on UTF-8 bytes, not character count — a document with many multi-byte characters can exceed the limit even if it appears short. Measure with new TextEncoder().encode(text).length before calling the API. For documents over the limit, split on sentence boundaries and call DetectEntities per chunk, then merge results with offset adjustment.

Async batch inference without an endpoint

import {
  StartEntitiesDetectionJobCommand,
  DescribeEntitiesDetectionJobCommand,
} from '@aws-sdk/client-comprehend';

// No endpoint required — submit an async job pointing at S3 input/output
const { JobId } = await client.send(new StartEntitiesDetectionJobCommand({
  JobName: `entity-batch-${Date.now()}`,
  DataAccessRoleArn: 'arn:aws:iam::123456789012:role/ComprehendDataAccessRole',
  LanguageCode: LanguageCode.EN,
  EntityRecognizerArn: entityRecognizerArn,
  InputDataConfig: {
    S3Uri: 's3://my-bucket/batch-input/',
    InputFormat: 'ONE_DOC_PER_LINE',  // or ONE_DOC_PER_FILE
  },
  OutputDataConfig: {
    S3Uri: 's3://my-bucket/batch-output/',
  },
}));

// Poll until COMPLETED (typically minutes to hours depending on volume)
async function waitForBatchJob(jobId: string): Promise<string> {
  while (true) {
    const { EntitiesDetectionJobProperties: props } = await client.send(
      new DescribeEntitiesDetectionJobCommand({ JobId: jobId })
    );
    const status = props?.JobStatus;
    if (status === 'COMPLETED') return props.OutputDataConfig?.S3Uri ?? '';
    if (status === 'FAILED') throw new Error(`Job failed: ${props?.Message}`);
    await new Promise(r => setTimeout(r, 30_000));
  }
}

Batch job output is written to <OutputS3Uri>/<AccountId>-NER-<JobId>/output/output.tar.gz. The output is a JSON Lines file matching the input structure with an added Entities array per document. Batch jobs are billed per character processed (~$0.0001/unit where 1 unit = 100 characters) with no endpoint charge — this is significantly cheaper than real-time endpoints for infrequent large-volume processing.

Confidence threshold design

The Score field on each returned entity represents the model's confidence that the span is the declared entity type. There is no universally correct threshold — the right value depends on your downstream use case and the relative cost of false positives vs false negatives:

For AliveMCP's use case (extracting MCP server slugs from issue reports to automatically link them to status pages), a threshold of 0.85 works well: high enough to avoid linking random strings to wrong servers, low enough to catch novel server slug formats the model has partial confidence in.

Tune threshold by evaluating your recognizer's per-score distribution on your test set: plot precision and recall at each 0.05 score increment, then pick the threshold at the knee of the precision-recall curve for your use case's tolerance.

Versioning and model lifecycle

Each CreateEntityRecognizerCommand call creates a new version under the same recognizer name. The EntityRecognizerArn returned is version-specific (contains /version/<name>). Endpoint deployment always points at a specific version ARN, not the recognizer name — you can canary test a new version by creating a second endpoint and routing a fraction of traffic to it before deleting the old endpoint.

List all versions with ListEntityRecognizerSummariesCommand (filter by RecognizerName). Delete old versions with DeleteEntityRecognizerCommand — you cannot delete a version that has a deployed endpoint. Endpoints must be deleted first (DeleteEndpointCommand), then the version can be deleted.