Python: feat: add Amazon Bedrock Knowledge Base tool and context provider - #8173
Python: feat: add Amazon Bedrock Knowledge Base tool and context provider#8173Vidyadhar Pogul (PVidyadhar) wants to merge 10 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.
🟡 Changes recommended
Agentic retrieval performs unused response generation, the IAM examples deny agentic calls, and provider behavior contradicts the PR description.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (2)
Previously missed (2) — in code that hasn't changed since the last review.
python/packages/bedrock/BEDROCK_MANAGED_KB.md:59
bedrock:AgenticRetrieveStreamdoes not support resource-level permissions, so this KB-scoped policy will deny the agentic request and force the tool down its fallback path. Grant that action in a separate statement withResource: "*", while retaining the KB ARN forbedrock:Retrieve.
python/packages/bedrock/samples/README.md:35bedrock:AgenticRetrieveStreamis not resource-scopable, so granting it only on a knowledge-base ARN produces an implicit deny. Split it into a separate statement withResource: "*"; keepbedrock:Retrievescoped to the KB ARN.
- Files reviewed: 13/14 changed files
- Comments generated: 2
- Review effort level: Balanced
| context.extend_messages( | ||
| self.source_id, | ||
| [Message(role="user", contents=[f"{self.context_prompt}\n\n{retrieved_context}"])], | ||
| ) |
GitHub Actions
Inactive
Addresses third Copilot review on microsoft#8173: - AgenticRetrieveStream defaults to generating a response (verified: 331 streamed responseEvents when omitted vs 0 with generateResponse=False). The tool only formats retrieval passages and discards generation, so pass generateResponse=False to avoid unnecessary model generation latency/cost. - Added a test asserting generateResponse=False is sent. - PR description updated separately to match the actual implementation (user-role injection + serializer coalescing, not extend_instructions). 68 unit tests pass; verified generateResponse behavior against live API.
GitHub Actions
Inactive
|
Thanks Copilot — both addressed:
68 unit tests pass; lint clean. |
There was a problem hiding this comment.
🟡 Changes recommended
IAM guidance, citation handling, public namespace exports, serializer documentation, and formatting need correction.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (5)
Previously missed (4) — in code that hasn't changed since the last review.
python/packages/bedrock/agent_framework_bedrock/init.py:8
- These package-level exports are not wired into the repository's documented public connector namespace.
agent_framework.amazon.__init__.pyand its.pyicurrently omit both names, sofrom agent_framework.amazon import BedrockKnowledgeBaseToolfails and type checkers cannot discover either API. Add both lazy runtime mappings and stub/__all__exports, consistent with the existing Bedrock classes there.
python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:45 - The fallback currently drops citations for valid Bedrock
Retrievelocation variants: the supported response union also includes Google Drive, OneDrive, Kendra, Salesforce, and SQL. Those results are still returned, but both the tool and provider render an empty source. Handle every location variant supported by the new SDK floor.
python/packages/bedrock/BEDROCK_MANAGED_KB.md:61 bedrock:AgenticRetrieveStreamhas no resource-level IAM type, so scoping it to a Knowledge Base ARN implicitly denies the agentic call and makes the documented tool silently fall back to standard retrieval. Keepbedrock:Retrievescoped to the KB ARN, but grantAgenticRetrieveStreamin a separate statement withResource: "*".
python/packages/bedrock/samples/README.md:37bedrock:AgenticRetrieveStreamis not resource-scopable, so this sample policy denies that action and the default tool never performs agentic retrieval. Split it into a statement withResource: "*", while retaining the Knowledge Base ARN restriction forbedrock:Retrieve.
python/packages/bedrock/agent_framework_bedrock/_chat_client.py:505
- This changes provider serialization but the required function-calling scenario matrix has not been updated.
python/AGENTS.md:60-67requires every provider-serialization change to updatedocs/specs/004-python-function-calling-loop.mdand its scenario-to-test mapping; add the adjacent-user invariant and the new Bedrock regression tests there.
# Coalesce adjacent genuine user-role turns only. Context providers
# (e.g. the Bedrock Knowledge Base provider) inject retrieved passages as
# separate user messages that would otherwise sit next to the real user
# input and violate Bedrock's role-alternation requirement. We restrict
- Files reviewed: 13/14 changed files
- Comments generated: 1
- Review effort level: Balanced
|
please use the default PR template Vidyadhar Pogul (@PVidyadhar) |
…ft#8173 - Revert the _prepare_bedrock_messages coalescing change entirely (now matches mainline). Provider instead relies on user-role injection like the repo's azure-cosmos-memory provider; consecutive user turns are tolerated by Bedrock Converse (verified E2E). This keeps us out of the protected function-calling serialization area (python/AGENTS.md), which requires core-team sign-off and spec updates for external contributors. - Wire BedrockKnowledgeBaseTool/Provider into the public agent_framework.amazon namespace (__init__.py lazy map + __init__.pyi stub + __all__), so 'from agent_framework.amazon import BedrockKnowledgeBaseTool' works and type checkers can discover both APIs. - Handle all Retrieve location variants in _get_source_uri (Google Drive, OneDrive, Salesforce, Kendra, SQL) per the boto3 >=1.43.32 schema, so citations are not dropped for those source types. - IAM examples: split bedrock:AgenticRetrieveStream into its own statement with Resource '*' (it has no resource-level type; scoping to a KB ARN implicitly denies it and forces silent fallback). bedrock:Retrieve stays KB-ARN scoped. - Strengthen default context prompt to frame passages as untrusted reference data (prompt-injection hardening). - Ruff-format/import-sort the changed test and source files. 65 unit tests pass; source + test lint clean; E2E verified against live KB.
GitHub Actions
Active
GitHub Actions
Active
|
Eduard van Valkenburg (@eavanvalkenburg) done — reformatted the PR description to follow the default template (Motivation & Context / Description & Review Guide / Related Issue / Contribution Checklist). Also addressed the latest Copilot review (db70b32):
65 unit tests pass; lint clean; E2E verified against a live KB. |
There was a problem hiding this comment.
🟡 Changes recommended
The provider can produce consecutive user messages that Bedrock Converse rejects, breaking the documented integration flow.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 13/14 changed files
- Comments generated: 2
- Review effort level: Balanced
Per python/samples/SAMPLE_GUIDELINES.md, provider-specific samples belong under python/samples/02-agents/providers/<provider>/. Move bedrock_kb_tool.py and bedrock_kb_context_provider.py from packages/bedrock/samples/ into python/samples/02-agents/providers/amazon/ alongside the existing bedrock_chat_client.py, so they are discoverable with the other Bedrock samples. - Align samples with the canonical convention: import from agent_framework.amazon, load_dotenv(), BedrockChatClient() with BEDROCK_* env vars. - Fold the KB usage guidance + IAM policy into the amazon provider README. - Point the package README at the relocated KB samples. Addresses Copilot review comment on sample location. 65 unit tests pass; samples ruff clean + formatted.
GitHub Actions
Active
There was a problem hiding this comment.
🟡 Changes recommended
The samples ignore the documented region setting, and two usage examples call an unsupported client constructor.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (3)
Previously missed (3) — in code that hasn't changed since the last review.
python/packages/bedrock/BEDROCK_MANAGED_KB.md:17
- The documented usage raises
TypeErrorbecauseBedrockChatClientdoes not acceptoptions; additionally, the chat option key ismodel, notmodel_id. Use the client's supportedmodelparameter so this copyable example works.
python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:78 - This usage example cannot run:
BedrockChatClient.__init__has nooptionsparameter, andBedrockChatOptionsusesmodel, notmodel_id. Construct the client with its supportedmodelargument instead.
python/packages/bedrock/tests/test_bedrock_knowledge_base.py:271 - This comment is inaccurate:
_prepare_bedrock_messages()appends each message unchanged and does not coalesce adjacent roles. Since the current behavior intentionally relies on Bedrock Converse accepting consecutive user messages, describe that behavior instead of documenting nonexistent coalescing.
- Files reviewed: 13/14 changed files
- Comments generated: 2
- Review effort level: Balanced
| # 1. Create the Knowledge Base context provider — subclasses ContextProvider. | ||
| kb_provider = BedrockKnowledgeBaseProvider( | ||
| knowledge_base_id="YOUR_KB_ID", # Replace with your managed KB ID | ||
| region_name="us-west-2", | ||
| number_of_results=3, | ||
| min_score=0.3, # Only include results above this relevance threshold | ||
| source_id="company-docs", # Unique ID for this context source | ||
| ) | ||
|
|
||
| # 2. Create an agent with the context provider — context is injected on every run. | ||
| agent = Agent( | ||
| client=BedrockChatClient(), |
There was a problem hiding this comment.
Good catch — fixed in 46fb12e. The sample advertised BEDROCK_REGION but hardcoded the KB client to us-west-2, so a KB configured via BEDROCK_REGION in another region would have been queried in the wrong region while BedrockChatClient used the configured one. Both KB samples now set region_name=os.environ.get("BEDROCK_REGION", "us-east-1"), matching BedrockChatClient and the documented env var.
| # 1. Create the Knowledge Base tool — subclasses FunctionTool, pass directly to Agent. | ||
| kb_tool = BedrockKnowledgeBaseTool( | ||
| knowledge_base_id="YOUR_KB_ID", # Replace with your managed KB ID | ||
| region_name="us-west-2", | ||
| number_of_results=5, | ||
| use_agentic_retrieval=True, # Uses query decomposition + managed reranking | ||
| ) | ||
|
|
||
| # 2. Create an agent with the KB tool — the agent calls it when it needs context. | ||
| agent = Agent( | ||
| client=BedrockChatClient(), |
There was a problem hiding this comment.
Good catch — fixed in 46fb12e. The sample advertised BEDROCK_REGION but hardcoded the KB client to us-west-2, so a KB configured via BEDROCK_REGION in another region would have been queried in the wrong region while BedrockChatClient used the configured one. Both KB samples now set region_name=os.environ.get("BEDROCK_REGION", "us-east-1"), matching BedrockChatClient and the documented env var.
Both KB samples advertised BEDROCK_REGION in their docstrings but hardcoded the KB client to us-west-2, so a KB configured via BEDROCK_REGION in another region would be queried in the wrong region (while BedrockChatClient used the configured region). Read region_name from os.environ['BEDROCK_REGION'] (default us-east-1), matching BedrockChatClient and the documented env var. Addresses Copilot review comments on both amazon KB samples. Samples ruff clean.
GitHub Actions
Active
The comment referenced _prepare_bedrock_messages coalescing, which was reverted. Retrieved content stays in the untrusted user channel (matches azure-cosmos-memory convention); Bedrock Converse tolerates consecutive user turns (verified E2E), so no coalescing is involved. Comment-only change.
GitHub Actions
Active
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