MCP Server Bedrock Knowledge Base — retrieveAndGenerate, citations, chunking, S3 sources

Retrieve vs RetrieveAndGenerate: choosing the right API for your MCP tool

The two Knowledge Base APIs serve different positions in the RAG pipeline. Retrieve runs only the retrieval step — it embeds your query, searches the vector store, and returns source chunks with similarity scores. RetrieveAndGenerate runs the full pipeline — retrieval plus model synthesis — returning a generated answer with citations. The choice determines where the RAG logic lives: in Bedrock (for RetrieveAndGenerate) or in your MCP server code (for Retrieve).

Use Retrieve when: you want to compose the retrieved chunks into a prompt yourself before calling a model (for example, to interleave retrieved evidence with tool call results); you need to use a model not supported by RetrieveAndGenerate; or you need to filter, re-rank, or post-process the retrieved chunks before synthesis. Use RetrieveAndGenerate when the default synthesis behavior is sufficient and you want the citation tracking handled for you.

import {
  BedrockAgentRuntimeClient,
  RetrieveCommand,
  RetrieveAndGenerateCommand,
  type RetrieveCommandInput,
  type RetrieveAndGenerateCommandInput,
  type RetrievedReference,
} from '@aws-sdk/client-bedrock-agent-runtime';

const client = new BedrockAgentRuntimeClient({
  region: process.env.AWS_REGION ?? 'us-east-1',
});

const KNOWLEDGE_BASE_ID = process.env.BEDROCK_KB_ID!;

// --- Retrieve: pure vector search, returns chunks with scores ---
async function retrieveChunks(
  query: string,
  numberOfResults = 5,
): Promise> {
  const input: RetrieveCommandInput = {
    knowledgeBaseId: KNOWLEDGE_BASE_ID,
    retrievalQuery: { text: query },
    retrievalConfiguration: {
      vectorSearchConfiguration: {
        numberOfResults,
        overrideSearchType: 'HYBRID',  // SEMANTIC or HYBRID
      },
    },
  };

  const response = await client.send(new RetrieveCommand(input));

  return (response.retrievalResults ?? []).map(result => ({
    text:  result.content?.text ?? '',
    score: result.score ?? 0,
    s3Uri: result.location?.s3Location?.uri ?? '',
  }));
}

// --- RetrieveAndGenerate: full RAG pipeline in one call ---
interface RagResult {
  answer: string;
  citations: Array<{
    text: string;         // the generated text segment this citation supports
    sources: Array<{
      content: string;   // the retrieved chunk content
      s3Uri: string;     // S3 URI of the source document
    }>;
  }>;
  sessionId?: string;    // can be reused for multi-turn conversations
}

async function retrieveAndGenerate(
  query: string,
  sessionId?: string,    // pass for multi-turn; omit to start a new session
): Promise {
  // modelArn must be a supported model — NOT an inference profile ID
  // Supported: anthropic.claude-3-sonnet-20240229-v1:0,
  //            anthropic.claude-3-haiku-20240307-v1:0,
  //            amazon.titan-text-premier-v1:0
  const modelArn =
    `arn:aws:bedrock:${process.env.AWS_REGION ?? 'us-east-1'}` +
    `::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0`;

  const input: RetrieveAndGenerateCommandInput = {
    input: { text: query },
    retrieveAndGenerateConfiguration: {
      type: 'KNOWLEDGE_BASE',
      knowledgeBaseConfiguration: {
        knowledgeBaseId: KNOWLEDGE_BASE_ID,
        modelArn,
        retrievalConfiguration: {
          vectorSearchConfiguration: {
            numberOfResults: 5,
            overrideSearchType: 'HYBRID',
          },
        },
      },
    },
    sessionId,
  };

  const response = await client.send(new RetrieveAndGenerateCommand(input));

  // Extract the generated answer
  const answer = response.output?.text ?? '';

  // Extract citations — the structure is nested
  const citations = (response.citations ?? []).map(citation => ({
    text: citation.generatedResponsePart?.textResponsePart?.text ?? '',
    sources: (citation.retrievedReferences ?? []).map(ref => ({
      content: ref.content?.text ?? '',
      s3Uri:   ref.location?.s3Location?.uri ?? '',
    })),
  }));

  return {
    answer,
    citations,
    sessionId: response.sessionId,
  };
}

The overrideSearchType parameter on vectorSearchConfiguration has two values: SEMANTIC runs pure vector similarity search; HYBRID combines vector similarity with BM25 keyword search and merges results using a reciprocal rank fusion algorithm. Hybrid search generally outperforms pure semantic search on queries that contain specific entity names, product codes, or technical terms that may not be well-represented in the embedding space. It has slightly higher latency (10–30 ms in most cases) and is available only for OpenSearch Serverless vector stores, not Aurora PostgreSQL.

Citation extraction: retrievedReferences and S3 location URIs

The RetrieveAndGenerateCommand response includes a citations array that links segments of the generated answer back to the source document chunks that informed them. Properly extracting and returning these citations in your MCP tool response is what separates a hallucination-prone tool from a grounded, auditable one.

The citation structure has two levels. Each entry in citations corresponds to a segment of the generated text, identified by generatedResponsePart.textResponsePart.text (the text segment) and .span.start / .span.end (character offsets in the full answer). Under each citation is a retrievedReferences array — these are the source chunks from the vector store that informed that specific segment. Each reference has content.text (the raw chunk content) and location.s3Location.uri (the S3 URI of the source document).

// Full citation extraction for MCP tool output
interface ToolCitation {
  generatedSegment: string;
  spanStart: number;
  spanEnd: number;
  sources: Array<{
    chunkText: string;
    documentUri: string;
    // Derived from S3 URI — useful for display
    documentKey: string;
    documentBucket: string;
  }>;
}

function extractCitations(
  response: Awaited>>,
): ToolCitation[] {
  return (response.citations ?? []).map(citation => {
    const part = citation.generatedResponsePart?.textResponsePart;

    const sources = (citation.retrievedReferences ?? []).map(ref => {
      const uri = ref.location?.s3Location?.uri ?? '';
      // S3 URI format: s3://bucket-name/path/to/object.pdf
      const match = uri.match(/^s3:\/\/([^/]+)\/(.+)$/);

      return {
        chunkText:      ref.content?.text ?? '',
        documentUri:    uri,
        documentBucket: match?.[1] ?? '',
        documentKey:    match?.[2] ?? '',
      };
    });

    return {
      generatedSegment: part?.text ?? '',
      spanStart:        part?.span?.start ?? 0,
      spanEnd:          part?.span?.end ?? 0,
      sources,
    };
  });
}

// Format citations for MCP tool response text
function formatCitationsForMcp(answer: string, citations: ToolCitation[]): string {
  // Deduplicate source documents across all citations
  const sourcesMap = new Map();
  let refIndex = 1;

  for (const citation of citations) {
    for (const source of citation.sources) {
      if (source.documentUri && !sourcesMap.has(source.documentUri)) {
        sourcesMap.set(source.documentUri, `[${refIndex++}]`);
      }
    }
  }

  const refsSection = [...sourcesMap.entries()]
    .map(([uri, label]) => {
      const key = uri.replace(/^s3:\/\/[^/]+\//, '');  // strip bucket prefix
      return `${label} ${key}`;
    })
    .join('\n');

  return `${answer}\n\nSources:\n${refsSection}`;
}

One practical consideration: retrievedReferences can be empty for a citation segment if Bedrock's model generated that segment from its parametric knowledge rather than the retrieved context. Always check for empty retrievedReferences and surface this in the MCP tool output — a response with no source citations is a signal that the answer may be hallucinated or that the query did not retrieve relevant chunks. Setting numberOfResults higher (10–20) increases the chance of finding relevant context, at the cost of larger prompts and higher token usage.

Chunking strategies: fixed-size, hierarchical, and semantic

The chunking strategy is configured on the Knowledge Base data source and applies at sync time — when Bedrock indexes documents from S3 into the vector store. Changing the chunking strategy requires deleting all data source sync jobs and re-syncing, which for large document sets can take hours and causes a period of degraded retrieval quality during re-indexing. Choose the right strategy before production rollout.

Fixed-size chunking splits documents into overlapping windows of a fixed character or token count. The window size (default 300 tokens) and overlap percentage (default 20%) are configurable. It is the fastest to index, predictable in memory usage, and appropriate for short, uniform documents like FAQ entries, product descriptions, or structured records. The failure mode is semantic boundary crossings: a chunk might start mid-sentence or split a table across two chunks, degrading retrieval quality for the second chunk.

Hierarchical chunking creates a parent-child chunk structure: large parent chunks (e.g. 1500 tokens) for broad context, with smaller child chunks (e.g. 300 tokens) used for retrieval. When a child chunk is retrieved, the corresponding parent chunk is injected into the synthesis context, giving the model broader context than the small retrieval hit alone would provide. This significantly improves answer quality for documents with multi-paragraph arguments, technical explanations, or narrative structure. It is 2–3x slower to index than fixed-size and consumes more storage in the vector store.

Semantic chunking uses a separate embedding model to identify natural breakpoints in the text — paragraph boundaries, topic shifts — and splits there instead of at fixed windows. It produces the most semantically coherent chunks and achieves the best retrieval recall on complex documents. The indexing cost is highest: it requires embedding calls for boundary detection, making it 5–10x slower than fixed-size for large document sets. Use it for unstructured documents like research papers, legal contracts, or support transcripts where the semantic structure is critical and sync frequency is low.

import {
  BedrockAgentClient,
  CreateDataSourceCommand,
  type ChunkingConfiguration,
} from '@aws-sdk/client-bedrock-agent';

const agentMgmt = new BedrockAgentClient({ region: 'us-east-1' });

// Fixed-size chunking configuration
const fixedSizeChunking: ChunkingConfiguration = {
  chunkingStrategy: 'FIXED_SIZE',
  fixedSizeChunkingConfiguration: {
    maxTokens: 300,
    overlapPercentage: 20,
  },
};

// Hierarchical chunking configuration
const hierarchicalChunking: ChunkingConfiguration = {
  chunkingStrategy: 'HIERARCHICAL',
  hierarchicalChunkingConfiguration: {
    levelConfigurations: [
      { maxTokens: 1500 },  // parent chunk size
      { maxTokens: 300 },   // child chunk size
    ],
    overlapTokens: 60,
  },
};

// Semantic chunking configuration
const semanticChunking: ChunkingConfiguration = {
  chunkingStrategy: 'SEMANTIC',
  semanticChunkingConfiguration: {
    maxTokens: 300,
    bufferSize: 0,           // number of surrounding sentences to include for boundary detection
    breakpointPercentileThreshold: 95,  // higher = fewer, larger chunks
  },
};

// Create a data source with hierarchical chunking
async function createS3DataSource(
  knowledgeBaseId: string,
  s3BucketArn: string,
  chunkingConfig: ChunkingConfiguration,
) {
  const { dataSource } = await agentMgmt.send(new CreateDataSourceCommand({
    knowledgeBaseId,
    name: 'production-docs',
    dataSourceConfiguration: {
      type: 'S3',
      s3Configuration: {
        bucketArn: s3BucketArn,
        inclusionPrefixes: ['docs/', 'kb/'],  // only index objects under these prefixes
      },
    },
    vectorIngestionConfiguration: {
      chunkingConfiguration: chunkingConfig,
    },
  }));
  return dataSource?.dataSourceId;
}

S3 data sources: metadata files and supported formats

Bedrock Knowledge Base indexes S3 objects during sync jobs. For each S3 object, you can provide a sidecar metadata file that defines structured attributes used for RetrievalFilter expressions. The metadata file must be placed at the same S3 path as the document, with .metadata.json appended to the object key.

Supported document formats for indexing: PDF, Markdown (.md), DOCX, HTML, CSV, TXT, and XLSX. Files in unsupported formats are silently skipped during sync — there is no error; the sync job completes successfully with those files absent from the index. Always verify your target file types are in the supported list before designing the document corpus, and check the sync job statistics (numberOfDocumentsScanned, numberOfNewDocumentsIndexed, numberOfDocumentsDeleted, numberOfDocumentsFailed) after each sync to catch format-related exclusions.

// S3 metadata file format
// Object key:          docs/2026/q3-earnings-report.pdf
// Metadata file key:   docs/2026/q3-earnings-report.pdf.metadata.json
// The metadata file contains the metadataAttributes object:

const metadataFileContent = {
  metadataAttributes: {
    // String attributes
    department:   { value: 'finance',       type: 'STRING' },
    docType:      { value: 'earnings',       type: 'STRING' },
    language:     { value: 'en',             type: 'STRING' },
    // Number attributes (for range queries)
    fiscalYear:   { value: 2026,             type: 'NUMBER' },
    fiscalQuarter:{ value: 3,                type: 'NUMBER' },
    // Boolean attributes
    public:       { value: true,             type: 'BOOLEAN' },
    // String list attributes (for IN-style filtering)
    regions:      { value: ['us', 'eu'],     type: 'STRING_LIST' },
  },
};

// Upload the metadata file alongside the document
import { S3Client, PutObjectCommand } from '@aws-sdk/client-s3';
const s3 = new S3Client({ region: 'us-east-1' });

async function uploadDocumentWithMetadata(
  bucket: string,
  objectKey: string,
  documentBody: Buffer,
  metadata: Record,
) {
  // Upload document
  await s3.send(new PutObjectCommand({
    Bucket: bucket,
    Key: objectKey,
    Body: documentBody,
  }));

  // Upload metadata sidecar
  await s3.send(new PutObjectCommand({
    Bucket: bucket,
    Key: `${objectKey}.metadata.json`,
    Body: JSON.stringify({ metadataAttributes: metadata }),
    ContentType: 'application/json',
  }));
}

Metadata attributes are indexed at sync time. If you add metadata to an existing S3 object after the initial sync, you must trigger a new sync job for the metadata to be available for filtering. Bedrock does not watch for metadata file changes between syncs. Sync jobs can be triggered manually via StartIngestionJobCommand or scheduled via Amazon EventBridge. For frequently updated knowledge bases, schedule sync jobs in off-peak hours and monitor sync job duration — hierarchical and semantic chunking can make sync jobs for large corpora take 30+ minutes.

RetrievalFilter: attribute-based pre-filtering before vector search

A RetrievalFilter constrains the vector search to only consider document chunks that match the filter expression, evaluated against the metadata attributes indexed from the .metadata.json sidecar files. The filter is applied before the embedding similarity search — it narrows the candidate set, not the ranked results. This is the correct tool for multi-tenant knowledge bases (filter by tenantId), versioned documentation (filter by productVersion), or time-bounded queries (filter by publishedYear greaterThan 2024).

The filter DSL supports atomic operators (equals, notEquals, greaterThan, greaterThanOrEquals, lessThan, lessThanOrEquals, in, notIn, startsWith, listContains) and logical operators (andAll, orAll, notFilter). Logical operators take an array of filter expressions, enabling arbitrary nesting.

import type { RetrievalFilter } from '@aws-sdk/client-bedrock-agent-runtime';

// Simple equality filter — only retrieve chunks from the 'finance' department
const departmentFilter: RetrievalFilter = {
  equals: {
    key: 'department',
    value: 'finance',
  },
};

// Range filter — only documents from fiscal year 2025 or later
const recencyFilter: RetrievalFilter = {
  greaterThanOrEquals: {
    key: 'fiscalYear',
    value: 2025,
  },
};

// Compound filter — finance department AND (fiscal year >= 2025 OR document is public)
const compoundFilter: RetrievalFilter = {
  andAll: [
    {
      equals: { key: 'department', value: 'finance' },
    },
    {
      orAll: [
        { greaterThanOrEquals: { key: 'fiscalYear', value: 2025 } },
        { equals: { key: 'public', value: true } },
      ],
    },
  ],
};

// Using a filter in RetrieveAndGenerate
async function ragWithFilter(
  query: string,
  tenantId: string,
  productVersion: string,
): Promise {
  const tenantFilter: RetrievalFilter = {
    andAll: [
      { equals: { key: 'tenantId',        value: tenantId } },
      { equals: { key: 'productVersion',  value: productVersion } },
    ],
  };

  const modelArn =
    `arn:aws:bedrock:${process.env.AWS_REGION ?? 'us-east-1'}` +
    `::foundation-model/anthropic.claude-3-haiku-20240307-v1:0`;

  const response = await client.send(new RetrieveAndGenerateCommand({
    input: { text: query },
    retrieveAndGenerateConfiguration: {
      type: 'KNOWLEDGE_BASE',
      knowledgeBaseConfiguration: {
        knowledgeBaseId: KNOWLEDGE_BASE_ID,
        modelArn,
        retrievalConfiguration: {
          vectorSearchConfiguration: {
            numberOfResults: 10,
            overrideSearchType: 'HYBRID',
            filter: tenantFilter,            // applied before vector similarity scoring
          },
        },
      },
    },
  }));

  return {
    answer:     response.output?.text ?? '',
    citations:  extractCitations(response),
    sessionId:  response.sessionId,
  };
}

// Filter builder helper for MCP tool that accepts filter params
function buildRetrievalFilter(
  params: Record,
): RetrievalFilter | undefined {
  const conditions: RetrievalFilter[] = [];

  for (const [key, value] of Object.entries(params)) {
    if (value === undefined || value === null) continue;
    if (typeof value === 'string' || typeof value === 'boolean') {
      conditions.push({ equals: { key, value } });
    } else if (typeof value === 'number') {
      conditions.push({ equals: { key, value } });
    } else if (Array.isArray(value)) {
      conditions.push({ in: { key, value } });
    }
  }

  if (conditions.length === 0) return undefined;
  if (conditions.length === 1) return conditions[0];
  return { andAll: conditions };
}

One important constraint: RetrievalFilter can only reference attributes that exist in the indexed metadata. If a document was synced before the metadata file was added (or if the metadata file is missing), those chunks will not match any attribute filter and will be excluded from results — this can cause partial or empty retrieval for documents with incomplete metadata. Validate metadata coverage as part of your sync pipeline by checking the count of indexed documents with specific attributes against the total document count.

Common failure modes