Skip to content

feat(bedrock-kb-retrieval-mcp-server): expose document metadata and add metadata_filter to QueryKnowledgeBases - #4330

Open
felipedbene wants to merge 3 commits into
awslabs:mainfrom
felipedbene:feat/bedrock-kb-metadata-filter
Open

feat(bedrock-kb-retrieval-mcp-server): expose document metadata and add metadata_filter to QueryKnowledgeBases#4330
felipedbene wants to merge 3 commits into
awslabs:mainfrom
felipedbene:feat/bedrock-kb-metadata-filter

Conversation

@felipedbene

@felipedbene felipedbene commented Jul 23, 2026

Copy link
Copy Markdown

Closes #2935

Description

Two backward-compatible changes to src/bedrock-kb-retrieval-mcp-server, as discussed in #2935:

  1. Expose document metadata in results. The Bedrock Retrieve API returns a metadata object per result (system x-amz-bedrock-kb-* attributes plus any custom attributes from .metadata.json sidecars), but knowledgebases/retrieval.py previously built the emitted document from content, location, and score only. Each emitted document now includes 'metadata': result.get('metadata', {}). Existing callers just see one additional key.

  2. Add optional metadata_filter parameter to QueryKnowledgeBases. A raw passthrough of the documented RetrievalFilter schema into retrievalConfiguration.vectorSearchConfiguration.filter. This imposes no schema opinions and supports any filter expression the Bedrock API supports (equals, notEquals, greaterThan, lessThan, in, startsWith, andAll, orAll, etc.). The existing data_source_ids filter logic was refactored so that when both are provided they are composed via andAll; a single condition stays bare, preserving the exact request shape emitted today for data_source_ids-only calls.

    The composition respects the Retrieve API's documented one-level filter embedding limit: a top-level andAll filter has the data-source condition merged into its member list (AND is associative, so this preserves both semantics and depth); any other filter is wrapped together with the data-source condition in a new andAll. The one unresolvable corner — a top-level orAll that already contains embedded filter groups, combined with data_source_ids — raises a ValueError with an actionable message rather than sending a request the API would reject. An empty metadata_filter object also raises an actionable ValueError (the API rejects an empty RetrievalFilter tagged union, so passing it through would only produce an opaque botocore ParamValidationError), and a vacuous {'andAll': []} merged with the data-source condition degrades to the bare condition to respect the API's two-member minimum for andAll.

Validation

  • Validated against a real knowledge base with ~26k documents (metadata present in results; metadata_filter alone; data_source_ids alone; both combined under andAll).
  • New unit tests covering: metadata present in output, metadata defaulting to {} when absent from the API response, filter passthrough alone (no andAll wrapper), empty-object filter rejected with an actionable ValueError, data_source_ids alone (unchanged request shape), simple filter + data_source_ids wrapped in andAll, top-level andAll merged in place (depth preserved), vacuous {'andAll': []} degrading to the bare data-source condition, flat orAll wrapped correctly, and the nested-orAll + data_source_ids ValueError.
  • Full test suite passes (49 tests), ruff check / ruff format clean, pyright 0 errors.

Type of change

  • New feature (non-breaking change which adds functionality)

Checklist

  • I have performed a self-review of this change
  • Changes have been tested
  • Changes are documented (server README tool docs, tool docstring, CHANGELOG)
  • Changes generate no new warnings
    By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of the project license.

…dd metadata_filter to QueryKnowledgeBases

- Include the Bedrock Retrieve API's per-result metadata object in each
  document emitted by query_knowledge_base (previously only content,
  location, and score were returned)
- Add optional metadata_filter parameter to query_knowledge_base and the
  QueryKnowledgeBases tool: a raw passthrough of the Bedrock
  RetrievalFilter schema into
  retrievalConfiguration.vectorSearchConfiguration.filter
- Compose metadata_filter with the existing data_source_ids filter via
  andAll when both are provided; a single condition stays bare

Closes awslabs#2935
@github-actions

Copy link
Copy Markdown
Contributor

This pull request is now marked as stale because it hasn't seen activity for a while. Add a comment or it will be closed soon. If you wish to exclude this issue from being marked as stale, add the "backlog" label.

@github-actions github-actions Bot added the stale These are items that have been around for a long time without progress label Aug 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

stale These are items that have been around for a long time without progress

Projects

Status: To triage

Development

Successfully merging this pull request may close these issues.

bedrock-kb-retrieval-mcp-server: expose metadata filtering in QueryKnowledgeBases tool

1 participant