Repository navigation
Python: feat: add Amazon Bedrock Knowledge Base tool and context provider - #8173
Vidyadhar Pogul (PVidyadhar) wants to merge 21 commits into
Conversation
- Created BedrockKnowledgeBaseTool with async run() + get_tool_definition() - Created BedrockKnowledgeBaseProvider (ContextProvider subclass) with before_run() - Two integration points: standalone tool + automatic context injection - Supports managed search and agentic retrieval with fallback - Unit tests included - Added BEDROCK_MANAGED_KB.md design doc
Addresses reviewer feedback (@moonbox3): when the provider is used with BedrockChatClient, injecting retrieved context as a separate user message produced two consecutive user turns in _prepare_bedrock_messages (which does not coalesce same-role messages). Route the retrieved context through extend_instructions() so it lands in Bedrock's system field, separate from the conversation array. This is model-agnostic and also avoids adding untrusted content as a system conversation message. - provider uses context.extend_instructions(self.source_id, ...) - removed unused Message import - updated tests to assert on context.instructions - 65 tests pass, verified E2E via agent.run() with BedrockChatClient + live KB
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
agent-framework-automation
Bot
added
documentation
GitHub Actions
Inactive
|
This PR continues from #7066, which could not be reopened after a force-push (GitHub rejected the reopen with a validation error). All prior review feedback from #7066 is carried over here. Notably addressing Evan Mattson (@moonbox3)'s comment from #7066 about consecutive cc Evan Mattson (@moonbox3) Eduard van Valkenburg (@eavanvalkenburg) — thanks for the earlier reviews on #7066. |
There was a problem hiding this comment.
🟡 Changes recommended
Retrieved content is elevated to system instructions, the lockfile is stale, and package guidance needs updating.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds Amazon Bedrock Knowledge Base retrieval as an agent tool and automatic context provider, replacing #7066.
Changes:
- Adds agentic retrieval with standard retrieval fallback.
- Adds automatic Knowledge Base context injection.
- Adds samples, tests, documentation, and newer AWS SDK requirements.
File summaries
| File | Description |
|---|---|
tests/test_bedrock_knowledge_base.py |
Tests tool and provider behavior. |
samples/README.md |
Documents sample patterns and permissions. |
samples/bedrock_kb_tool.py |
Demonstrates tool-based retrieval. |
samples/bedrock_kb_context_provider.py |
Demonstrates provider-based retrieval. |
samples/__init__.py |
Initializes the samples package. |
pyproject.toml |
Raises AWS SDK dependency floors. |
BEDROCK_MANAGED_KB.md |
Documents managed Knowledge Base support. |
_knowledge_base.py |
Implements the retrieval tool. |
_knowledge_base_provider.py |
Implements automatic context retrieval. |
__init__.py |
Exports the new public APIs. |
Review details
- Files reviewed: 10/10 changed files
- Comments generated: 3
- Review effort level: Balanced
💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.
| "boto3>=1.43.32,<2.0.0", | ||
| "botocore>=1.43.32,<2.0.0", |
| from ._knowledge_base import BedrockKnowledgeBaseTool | ||
| from ._knowledge_base_provider import BedrockKnowledgeBaseProvider |
- Keep retrieved KB passages as untrusted user-role context instead of elevating to system instructions (matches azure-cosmos-memory convention; avoids stored prompt-injection). Solve Bedrock role alternation by coalescing adjacent user-role messages in _prepare_bedrock_messages (assistant turns left untouched to preserve tool-use/tool-result pairing). - Regenerate python/uv.lock for the boto3/botocore >=1.43.32 floor. - Add BedrockKnowledgeBaseTool/Provider to bedrock AGENTS.md class list. - Tests: coalescing + no-coalesce-across-assistant cases; 67 pass. Verified E2E via agent.run() with BedrockChatClient + live KB.
GitHub Actions
Inactive
|
Thanks Copilot — addressed all three in 9f8ea75:
67 unit tests pass; lint clean. |
There was a problem hiding this comment.
🟡 Changes recommended
Agentic results are incorrectly formatted, valid source types are omitted, and the implementation contradicts the stated context-injection design.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (3)
Previously missed (3) — in code that hasn't changed since the last review.
python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:45
- This extractor omits valid Retrieve location variants, so citations are blank for Salesforce, Kendra, SQL, OneDrive, and Google Drive knowledge-base results even though the supported SDK response union includes them. Handle every location variant exposed by the dependency floor/current supported releases.
python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:153 - AgenticRetrieveStream synthesizes a generated answer by default, but this implementation discards both
generatedResponseandresponseEventand only formats retrieved items. That adds avoidable model latency and cost on every agentic lookup; disable response generation when requesting retrieval-only output.
python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py:139 - This provider does not implement the PR description's stated fix: the description says retrieved passages now use
context.extend_instructions()and that tests assertcontext.instructions, while this code and its test still add a user-role context message and rely on a new global serializer behavior. Please align the implementation/tests and description so the intended trust boundary and compatibility behavior are reviewable.
- Files reviewed: 13/14 changed files
- Comments generated: 2
- Review effort level: Balanced
Addresses second Copilot review on microsoft#8173: 1. AgenticRetrieveStream results use a different schema (content/metadata/ sourceRetriever) than standard Retrieve (score/location). Previously every agentic result was normalized to score 0 with a blank source. Now parse the source URI from metadata._source_uri and omit the score (managed reranking does not expose one); the formatter only renders a score when present. Updated the agentic test mock to the real SDK schema. 2. Restrict _prepare_bedrock_messages coalescing to messages whose ORIGINAL role is 'user', so tool-result turns (role='tool', which map to Bedrock 'user') are never merged into a preceding user text turn. This keeps function-call/tool-result serialization unchanged. Added a regression test for the tool-call/tool-result path. Verified E2E against live KB: agentic results show real source URLs and no fabricated scores. 68 unit tests pass.
GitHub Actions
Inactive
|
Thanks Copilot — both addressed in dfd0d2f:
68 unit tests pass; lint clean. |
There was a problem hiding this comment.
🔵 Needs a closer look
Audio and video retrieval results currently become blank passages because their non-text content types are not handled.
Review details
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:92
- Bedrock's
RetrievalResultContent.typealso includesAUDIOandVIDEO. Those payloads do not usecontent.text, so they currently fall through to line 94 and produce blank passages just like theIMAGEcase this branch fixes. Treat all unsupported binary media types explicitly (and update the helper's docstring/tests accordingly).
- Files reviewed: 13/14 changed files
- Comments generated: 0 new
- Review effort level: Balanced
Bedrock RetrievalResultContent.type is TEXT | IMAGE | AUDIO | VIDEO | ROW (verified against the botocore service model). AUDIO and VIDEO also carry byteContent rather than text, so they were falling through to the empty-text default and producing blank passages. Generalize the binary-media branch to IMAGE/AUDIO/VIDEO via _BINARY_MEDIA_CONTENT_TYPES, each rendered as a typed placeholder. Adds AUDIO/VIDEO tests.
|
Addressed the latest Copilot review (commit for AUDIO/VIDEO): confirmed against the botocore service model that |
GitHub Actions
Inactive
There was a problem hiding this comment.
🔵 Needs a closer look
The new clients bypass established Bedrock settings and use prohibited module-wide type-checking suppressions.
Review details
Suppressed comments (4)
Previously missed (2) — in code that hasn't changed since the last review.
python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:2
- The package coding standard explicitly forbids
# type: ignorein source modules (python/CODING_STANDARD.md:97-106), and this file-wide suppression disables checking for the new public API and all retrieval parsing. Please model the untyped boto3 surface withAny/a small protocol and use only rule-specific, line-level# pyright: ignore[...]suppressions where unavoidable; the legacy_chat_client.pysuppression should not be copied into new code.
This issue also appears on line 198 of the same file.
python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py:2
- The package coding standard explicitly forbids
# type: ignorein source modules (python/CODING_STANDARD.md:97-106). This file-wide suppression also hides mistakes in the new public provider API; please type the untyped boto3 boundary explicitly and use only rule-specific, line-level# pyright: ignore[...]suppressions where necessary rather than copying the legacy_chat_client.pyexception.
This issue also appears on line 88 of the same file.
python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:202
- Creating the global boto3 client directly bypasses the connector's established Bedrock settings/session path. Unlike
BedrockChatClient(_chat_client.py:290-312) and the embedding client, this tool ignoresBEDROCK_REGION,BEDROCK_ACCESS_KEY,BEDROCK_SECRET_KEY,BEDROCK_SESSION_TOKEN, and a reusable boto3 session. Consequently, a normally configuredBedrockChatClient()and default KB tool can target different regions or credentials. Resolve the same settings and expose the same region/credential/session options before constructing the agent-runtime client.
self._client = boto3.client(
"bedrock-agent-runtime",
region_name=self.region_name,
config=BotoConfig(user_agent_extra=f"{get_user_agent()} bedrock-kb"),
)
python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py:92
- This direct global boto3 client bypasses the Bedrock connector's settings/session behavior.
BedrockChatClientand the embedding client resolveBEDROCK_REGIONandBEDROCK_*credentials and accept a boto3 session, whereas this provider silently defaults to another region/credential chain. A provider attached to a configured chat client can therefore query the wrong region or fail authentication unless users duplicate configuration manually. Align this constructor with the existing Bedrock region/credential/session options.
self._client = boto3.client(
"bedrock-agent-runtime",
region_name=self.region_name,
config=BotoConfig(user_agent_extra=f"{get_user_agent()} bedrock-kb"),
)
- Files reviewed: 13/14 changed files
- Comments generated: 0 new
- Review effort level: Balanced
…l type:ignore Addresses review feedback on both KB files: - Remove the file-level '# type: ignore' (forbidden in source per CODING_STANDARD.md; the legacy _chat_client.py exception should not be copied). Type the untyped bedrock-agent-runtime boundary explicitly (annotate responses as dict[str, Any]) and use only targeted line-level '# pyright: ignore[reportUnknownMemberType]' on the two dynamic client calls. - Resolve region and credentials through the shared Bedrock settings/session path (BEDROCK_REGION / BEDROCK_ACCESS_KEY / BEDROCK_SECRET_KEY / BEDROCK_SESSION_TOKEN, .env files, and an optional boto3 Session) via a shared _build_kb_client() helper, matching BedrockChatClient and the embedding client. Both the tool and the provider now expose region/access_key/secret_key/session_token/boto3_session/env_file_* options, so a KB tool and a configured BedrockChatClient() no longer silently target different regions or credentials. Verified: ruff + 71 tests pass (1 integration skipped); E2E against a live managed KB confirms region resolves from BEDROCK_REGION for both the tool and the provider.
GitHub Actions
Inactive
|
Addressed the latest review (60b3361) — both points:
Verified: |
There was a problem hiding this comment.
🟡 Changes recommended
Explicit KB regions can be overridden by supplied sessions, and the new docstrings reverse the actual settings precedence.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 13/14 changed files
- Comments generated: 3
- Review effort level: Balanced
|
|
||
| return boto3_session.client( | ||
| "bedrock-agent-runtime", | ||
| region_name=boto3_session.region_name or resolved_region, |
| Region and credentials are resolved the same way as ``BedrockChatClient`` and | ||
| the embedding client — from these arguments, then the ``BEDROCK_*`` environment | ||
| variables (``BEDROCK_REGION``, ``BEDROCK_ACCESS_KEY``, ``BEDROCK_SECRET_KEY``, | ||
| ``BEDROCK_SESSION_TOKEN``), then an optional .env file — so a KB tool and a | ||
| configured ``BedrockChatClient()`` target the same region/credentials by default. |
| Region and credentials are resolved the same way as ``BedrockChatClient`` and | ||
| the embedding client — from these arguments, then the ``BEDROCK_*`` environment | ||
| variables, then an optional .env file — so this provider and a configured | ||
| ``BedrockChatClient()`` target the same region/credentials by default. |
| _BINARY_MEDIA_CONTENT_TYPES = frozenset({"IMAGE", "AUDIO", "VIDEO"}) | ||
|
|
||
|
|
||
| class _KnowledgeBaseSettings(TypedDict, total=False): |
There was a problem hiding this comment.
this does not need to be private, should be BedrockKnowledgeBaseSettings
There was a problem hiding this comment.
Done in 5defe1d — renamed to public BedrockKnowledgeBaseSettings and exported it.
| if col.get("columnName") or col.get("columnValue") | ||
| ] | ||
| return "\n".join(rendered) | ||
| if content_type in _BINARY_MEDIA_CONTENT_TYPES: |
There was a problem hiding this comment.
many models are multi-modal and we have Content.from_data and Content.from_uri (depending on data bytes or a URI pointing somewhere) that we can use for this
There was a problem hiding this comment.
Done in 5defe1d — image/audio/video passages are now injected as multi-modal Content (Content.from_uri with the source URI) via the provider's inject path, instead of a text placeholder. Text and SQL ROW passages stay as strings. Added a test asserting an IMAGE passage becomes a Content item.
| class _BedrockKBQueryInput(BaseModel): | ||
| """Input schema for the Bedrock Knowledge Base tool.""" | ||
|
|
||
| query: Annotated[str, Field(description="The search query to find relevant documents in the knowledge base.")] |
There was a problem hiding this comment.
creating a basemodel for this simple query only param definition is not great, just supply the json directly
There was a problem hiding this comment.
Done in an earlier commit (f5588c2) — dropped the _BedrockKBQueryInput BaseModel; the tool now supplies its single query param as an inline JSON schema.
|
|
||
| async def _retrieve(self, query: str) -> str: | ||
| """Retrieve and format context from the knowledge base.""" | ||
| passages = await asyncio.to_thread( |
There was a problem hiding this comment.
no agentic retrieval here?
There was a problem hiding this comment.
Added in 5defe1d — agentic retrieval (query decomposition + managed reranking, with standard fallback) is now shared via _retrieve_passages and used on both the inject and tool paths, controlled by use_agentic_retrieval (default True).
| logger = logging.getLogger("agent_framework.bedrock") | ||
|
|
||
|
|
||
| class BedrockKnowledgeBaseProvider(ContextProvider): |
There was a problem hiding this comment.
I think I would prefer if the user does everything through this context provider, where they can choose to:
- inject KB info (current behavior)
- add the KB search tool
- do both
Then we only have to have one public class that people work with and they can do the same thing.
There was a problem hiding this comment.
Restructured in 5defe1d. BedrockKnowledgeBaseProvider is now the single public class with a mode: inject (inject KB context), tool (expose the KB search tool), or both (default). before_run calls extend_messages and/or extend_tools accordingly. BedrockKnowledgeBaseTool is now internal (removed from __all__); the provider builds and exposes it for tool/both modes. Samples, README, AGENTS.md and tests updated; verified all three modes end-to-end against a live KB.
There was a problem hiding this comment.
this file will not do much...
There was a problem hiding this comment.
Removed in f5588c2 — the standalone BEDROCK_MANAGED_KB.md is gone; the class overview lives in AGENTS.md and the full agentic-retrieval IAM policy is in the sample README.
… standalone doc - Make the KB settings type public: _KnowledgeBaseSettings -> BedrockKnowledgeBaseSettings, exported from the package (review comment 1). - Drop the _BedrockKBQueryInput BaseModel; supply the tool's single-param input schema as an inline JSON schema instead (review comment 3). - Remove the standalone BEDROCK_MANAGED_KB.md; its class list is in AGENTS.md and the full IAM policy lives in the sample README (review comment 6).
GitHub Actions
Inactive
…de=inject|tool|both Per review: consolidate to a single public class. BedrockKnowledgeBaseProvider is now the entry point with mode='inject' (inject retrieved passages as context), 'tool' (expose a KB search tool the model can call), or 'both' (default). before_run() calls extend_messages and/or extend_tools accordingly. - BedrockKnowledgeBaseTool is now internal (removed from public __all__); the provider constructs and exposes it for tool/both modes. - Agentic retrieval (query decomposition + managed reranking, with standard fallback) is shared via _retrieve_passages and available on both the inject and tool paths. - Image/audio/video passages are injected as multi-modal Content (Content.from_uri) rather than text placeholders, for multi-modal models. - Settings type made public (BedrockKnowledgeBaseSettings); dropped the query BaseModel in favor of an inline JSON schema; removed the standalone BEDROCK_MANAGED_KB.md. - Updated tests (modes, tool-not-public, multimodal, agentic default), samples (tool sample uses mode='tool', context sample uses mode='inject'), README, and AGENTS.md. Verified: ruff + format clean, 77 tests pass (1 integration skipped); E2E against a live managed KB confirms all three modes and BEDROCK_REGION resolution.
GitHub Actions
Inactive
|
Thanks Eduard van Valkenburg (@eavanvalkenburg) — addressed all six comments, most significantly the consolidation to a single public class (5defe1d):
Samples now show |
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Agentic score filtering, multimodal serialization, and inconsistent public exports currently break documented behavior.
Get a fresh assessment by requesting another Copilot review.
Review effort: Balanced
Findings: 5
Open (9)
region_name(orBEDROCK_REGION) is ignored whenever a supplied boto3 session…python/uv.lockwith this dependency-floor change. Its…load_settings()resolves explicit…load_settings()resolves explicit…python/packages/bedrock/AGENTS.md:5-10…
| if p.score < self.min_score: | ||
| continue |
| "BedrockKnowledgeBaseProvider": ("agent_framework_bedrock", "agent-framework-bedrock"), | ||
| "BedrockKnowledgeBaseTool": ("agent_framework_bedrock", "agent-framework-bedrock"), |
|
|
||
| For Amazon Bedrock managed Knowledge Base retrieval, see: | ||
|
|
||
| - [`bedrock_kb_tool.py`](../../samples/02-agents/providers/amazon/bedrock_kb_tool.py) — `BedrockKnowledgeBaseTool` as a `FunctionTool` the agent calls on demand. |
|
Vidyadhar Pogul (@PVidyadhar) This branch currently conflicts with |
|
Vidyadhar Pogul (@PVidyadhar) Thanks for the substantial work on this. The branch still conflicts with |
|
Eduard van Valkenburg (@eavanvalkenburg) Sorry for missing your Sep 28 request. I've rebased onto current |
Motivation & Context
Enables Agent Framework agents to retrieve context from Amazon Bedrock Knowledge Bases, adding RAG capabilities on AWS's managed infrastructure without requiring users to run their own vector stores or embedding pipelines. This contributes the Amazon Bedrock Knowledge Base scenario to the Bedrock connector. Continues from #7066 (that PR could not be reopened via the UI after a force-push).
Description & Review Guide
BedrockKnowledgeBaseTool— subclassesFunctionToolwith agentic retrieval (AgenticRetrieveStream; query decomposition + managed reranking) and automatic fallback to standardRetrieve. Passable directly to any Agent or ChatClient.generateResponse=Falseis set so the tool returns passages only and the agent's own model generates the answer.BedrockKnowledgeBaseProvider— subclassesContextProvider; itsbefore_run()retrieves passages and injects them as an untrusted user-role message (same convention as theazure-cosmos-memoryprovider, avoiding elevation of retrieved content to system instructions).agent_framework.amazonnamespace.python/packages/bedrock/plus namespace exports. No change to shared serialization or the function-calling loop.FunctionTool/ContextProvidersubclassing conventions and the untrusted user-role injection inbefore_run().Related Issue
N/A — new feature. No existing open issue or PR (supersedes closed #7066).
Contribution Checklist