Repository navigation
Python: add Amazon Bedrock Knowledge Base context provider - #9782
Open
Vidyadhar Pogul (PVidyadhar) wants to merge 3 commits into
Open
Vidyadhar Pogul (PVidyadhar) wants to merge 3 commits into
Vidyadhar Pogul (PVidyadhar) wants to merge 3 commits into
Conversation
Adds BedrockKnowledgeBaseProvider, the single public entry point for using an Amazon Bedrock managed Knowledge Base with an agent. Its mode selects how the KB is used: 'inject' (retrieve passages and inject them as context), 'tool' (expose a KB search tool the model can call), or 'both' (default). - Agentic retrieval (AgenticRetrieveStream: query decomposition + managed reranking) with fallback to standard Retrieve, on both the inject and tool paths. - Image/audio/video results are injected as multi-modal Content; SQL ROW results are rendered as column/value text. - Region and credentials resolve via BEDROCK_* settings / boto3 session, matching BedrockChatClient and the embedding client. BedrockKnowledgeBaseSettings is public. - Retrieved content is injected as an untrusted user-role message. - Raises the boto3/botocore floor to 1.43.32 (first release with AgenticRetrieveStream). Rebased onto current main as a single commit. Also exports BedrockKnowledgeBaseSettings (instead of the now-internal tool) from the agent_framework.amazon lazy-import shim.
Vidyadhar Pogul (PVidyadhar)
temporarily deployed
to
github-app-auth
October 28, 2026 07:15 — with 
GitHub Actions
Inactive
GitHub Actions
Inactive
Vidyadhar Pogul (PVidyadhar)
requested review from
Darren Cohen (dargilco),
Jose Alvarez (jpalvarezl),
sophia-ramsey and
westey (westey-m)
as code owners
October 28, 2026 07:15
5 tasks done
Vidyadhar Pogul (PVidyadhar)
temporarily deployed
to
github-app-auth
October 28, 2026 07:15 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
agent-framework-automation
Bot
added
documentation
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Multiple knowledge-base providers generate duplicate tool names, causing tool-map validation to raise ValueError and preventing searches when providers are attached in tool or default both mode.
4 open issues
What changed in this PR
Adds Amazon Bedrock Knowledge Base retrieval to the Python connector through a single context provider, supporting automatic context injection, model-driven search, or both.
Changes:
- Adds shared agentic retrieval with standard retrieval fallback, passage formatting, and region/credential resolution.
- Exports the provider and settings through both public namespaces and raises the AWS SDK minimum versions.
- Adds unit tests, usage samples, and setup documentation.
| File | Description |
|---|---|
python/uv.lock |
Updates AWS SDK dependency constraints. |
python/samples/02-agents/providers/amazon/README.md |
Documents modes, samples, and retrieval permissions. |
python/samples/02-agents/providers/amazon/bedrock_kb_tool.py |
Demonstrates model-driven Knowledge Base search. |
python/samples/02-agents/providers/amazon/bedrock_kb_context_provider.py |
Demonstrates automatic context injection. |
python/packages/core/agent_framework/amazon/__init__.pyi |
Adds typed provider and settings exports. |
python/packages/core/agent_framework/amazon/__init__.py |
Adds lazy provider and settings exports. |
python/packages/bedrock/tests/test_bedrock_knowledge_base.py |
Tests retrieval, provider modes, filtering, and content handling. |
python/packages/bedrock/README.md |
Links Knowledge Base examples. |
python/packages/bedrock/pyproject.toml |
Raises boto3 and botocore minimums to 1.43.32. |
python/packages/bedrock/AGENTS.md |
Describes the new provider and settings. |
python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py |
Implements shared retrieval helpers and the internal search tool. |
python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py |
Implements injection, tool, and combined modes. |
python/packages/bedrock/agent_framework_bedrock/__init__.py |
Exports the public provider and settings. |
Review effort: Balanced
Comment on lines
+142
to
+143
| self._tool = BedrockKnowledgeBaseTool( | ||
| knowledge_base_id=self.knowledge_base_id, |
There was a problem hiding this comment.
Fixed in 8229f1e. The tool name now comes from source_id: the default source_id="bedrock-kb" still gives bedrock_knowledge_base, and any other value gives bedrock_knowledge_base_<sanitized source_id>, truncated to 64 characters. So several providers on one agent no longer collide. I also added tool_name / tool_description overrides, validated against the Bedrock tool-name pattern [a-zA-Z0-9_-]{1,64}. A test checks that two providers pass _ensure_unique_tool_names, and I ran a live check with two tool-mode providers on a single agent.
|
|
||
| return boto3_session.client( | ||
| "bedrock-agent-runtime", | ||
| region_name=boto3_session.region_name or resolved_region, |
There was a problem hiding this comment.
Addressed in 7ead64f (this thread is now outdated).
Comment on lines
+224
to
+225
| if p.score < self.min_score: | ||
| continue |
There was a problem hiding this comment.
Addressed in 7ead64f (this thread is now outdated).
| if content == f"[{media_type_name} content omitted]" and passage.source: | ||
| # e.g. IMAGE -> "image/*". Bedrock does not return an exact mime type here, so | ||
| # use the broad top-level type; multi-modal models resolve the concrete format. | ||
| return Content.from_uri(uri=passage.source, media_type=f"{media_type_name}/*") |
There was a problem hiding this comment.
Fixed in 8229f1e. Images now become Content.from_data(bytes, "image/<fmt>") with a concrete MIME type. That is the only image form BedrockChatClient serializes, so they reach Converse. Standard Retrieve data URIs and agentic byteContent + mimeType are both decoded. Formats Converse does not accept (anything other than gif/jpeg/png/webp) become a text placeholder that includes the source, so they are not silently dropped. Audio/video results use their transcription / summary text. Tests check that the image block reaches the Converse request.
…e and docs Addresses review findings from microsoft#8173 that were not yet handled: - Agentic retrieval exposes no numeric score; it was normalized to 0, so any positive min_score (the inject sample uses 0.3) filtered out every agentic passage and injected no context. Represent the missing score as None and apply min_score only to scored (standard Retrieve) passages. Adds a regression test. - A resolved region (argument / .env / BEDROCK_REGION) now takes precedence over a supplied boto3_session's default region, matching BedrockChatClient. Adds tests. - Docstrings now state load_settings' actual order (arguments, then .env file, then BEDROCK_* environment variables) and that before_run runs once per agent run. - Package README describes the provider's tool/inject modes instead of the internal tool.
Vidyadhar Pogul (PVidyadhar)
temporarily deployed
to
github-app-auth
October 28, 2026 07:21 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
- Derive the KB tool name from source_id (default stays bedrock_knowledge_base) so several providers on one agent no longer collide; add tool_name and tool_description overrides, validated against the Bedrock tool-name pattern. - Inject retrieved images as data Content with a concrete image MIME type (gif/jpeg/png/webp) so BedrockChatClient actually sends them; other formats become a text placeholder. Audio/video use their transcription/summary text. - Make the KB modules pass strict Pyright: typed views of the untyped bedrock-agent-runtime responses, non-private cross-module helpers, and direct boto3 imports (boto3 is a required dependency).
Vidyadhar Pogul (PVidyadhar)
temporarily deployed
to
github-app-auth
October 28, 2026 07:58 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
1 of 2 tasks
Vidyadhar Pogul (PVidyadhar)
temporarily deployed
to
github-app-auth
October 28, 2026 08:03 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
This branch was previously deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation & Context
Enables Agent Framework agents to retrieve context from Amazon Bedrock managed Knowledge Bases, adding RAG on AWS's managed infrastructure without users running their own vector stores or embedding pipelines. This contributes the Amazon Bedrock Knowledge Base scenario to the Bedrock connector.
Continues #8173, which could not be reopened after rebasing onto current
main(GitHub blocks reopening a closed PR once its branch has been force-pushed). Same branch, rebased with conflicts resolved, carrying forward everything addressed in #8173's review.Description & Review Guide
BedrockKnowledgeBaseProvider(subclassesContextProvider) is the single public entry point. Itsmodeselects how the KB is used:"inject"(before_run()retrieves passages and injects them as an untrusted user-role message viaextend_messages),"tool"(exposes a KB search tool the model can call on demand viaextend_tools), or"both"(default).AgenticRetrieveStream: query decomposition + managed reranking,generateResponse=False) with fallback to standardRetrieve, on both paths (use_agentic_retrieval, defaultTrue). Agentic results expose no numeric score, somin_scoreapplies only to scored (standardRetrieve) passages.Content(Content.from_datawith a concrete gif/jpeg/png/webp MIME type, soBedrockChatClientsends them to Converse). Other image formats become a text placeholder that names the source. Audio/video use their transcription/summary text, and SQLROWresults are rendered ascolumn: valuetext.source_id(defaultbedrock_knowledge_base), so several providers on one agent don't collide. It can be overridden withtool_name/tool_description..envfile, thenBEDROCK_*env vars; optionalboto3_session), with the same region precedence asBedrockChatClient.BedrockKnowledgeBaseSettingsis public.agent_framework_bedrockandagent_framework.amazon(BedrockKnowledgeBaseProvider,BedrockKnowledgeBaseSettings); the search tool is internal.boto3/botocorefloor to1.43.32, the first release withAgenticRetrieveStream;uv.lockregenerated.python/packages/bedrock/, two namespace exports, and the boto floor bump. No change to shared serialization or the function-calling loop.modedesign and its use ofextend_messages/extend_toolsinbefore_run().All review findings on #8173 are addressed here, including the Sep 21 design review (single public provider with
mode, agentic retrieval on the provider, multi-modalContent, public settings, inline JSON schema, standalone.mdremoved) and the later Copilot findings (agentic results vs.min_score, region precedence vs. a supplied session,load_settingsprecedence wording,agent_framework.amazonexports, README naming the internal tool). The Copilot review on this PR is addressed too: unique tool names across providers, and images reaching the Converse request.Validation:
ruff check/ruff format --checkclean (ruff 0.16.5 and 0.17.3). Strict Pyright reports 0 errors on the package source (pyright 1.1.414). The test/sample checkers (mypy, zuban, pyrefly, ty, pyright tests config) are clean on the bedrock tests and the amazon samples. Bedrock package tests pass against currentmain's core (1 integration test skipped without credentials). An end-to-end run against a live managed KB confirms all three modes, the sample configuration (min_score=0.3), two tool-mode providers on one agent, andBEDROCK_REGIONresolution.Related Issue
Fixes #9786
The design direction (single public provider with
mode) came out of Eduard van Valkenburg (@eavanvalkenburg)'s review of #8173. #9786 was filed to record the problem under the issue-first policy; I'm happy to adjust scope based on feedback there. Supersedes closed #8173 (and #7066 before it); there is no other open PR for this issue.AI Assistance
AI assistance details: implementation, unit tests, samples and documentation, addressing review feedback, and this PR description were prepared with an AI coding assistant. I reviewed all changes, ran the builds, linters and tests, and verified behavior end-to-end against a live Bedrock managed Knowledge Base.
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.