Python: feat: add Amazon Bedrock Knowledge Base tool and context provider

[Vidyadhar Pogul (PVidyadhar)]

feat: add Amazon Bedrock Knowledge Base tool and context provider

Continues from #7066 (that PR could not be reopened via the UI after a force-push). Same branch, all prior review feedback addressed.

Motivation & Context

Enables Agent Framework agents to retrieve context from Amazon Bedrock Knowledge Bases. Adds RAG capabilities using AWS's managed infrastructure without requiring users to manage their own vector stores or embedding pipelines.

Description & Review Guide

Major changes:

• BedrockKnowledgeBaseTool — subclasses FunctionTool with agentic retrieval (AgenticRetrieveStream, query decomposition + managed reranking) and automatic fallback to standard Retrieve. Can be passed directly to any Agent or ChatClient. Response generation is disabled (generateResponse=False) since the tool returns passages only and the agent's own model generates the answer.
• BedrockKnowledgeBaseProvider — subclasses ContextProvider with before_run() for automatic context injection on every agent invocation.
• A small serialization change in BedrockChatClient._prepare_bedrock_messages to coalesce adjacent genuine user-role turns (see below).

Impact: Additive feature + one scoped serializer change.

How the reviewer feedback was addressed

Role alternation with BedrockChatClient (Evan Mattson (@moonbox3), #7066): The provider injects retrieved passages as an untrusted user-role message (consistent with azure-cosmos-memory's convention of not elevating retrieved content to system instructions, which would open a stored prompt-injection path). To satisfy Bedrock's role-alternation requirement, _prepare_bedrock_messages now coalesces adjacent user-role turns. This is scoped to messages whose original role is user, so tool-result turns (role="tool", which map to the Bedrock user role) are never merged — function-call/tool-result serialization is unchanged. Regression tests cover both the coalescing and the tool-call/tool-result path.

Agentic result schema: AgenticRetrieveStream results expose content/metadata/sourceRetriever (no score/location). The source URI is read from metadata._source_uri, and no numeric score is fabricated for agentic results (managed reranking doesn't expose one); the formatter only renders a score when present.

Other Copilot items: regenerated python/uv.lock for the boto3/botocore >=1.43.32 floor; added the two KB classes to the package AGENTS.md; IAM examples wrapped as full policy documents; feature usage recorded via mark_feature_used(FeatureIndex.BEDROCK); logger name agent_framework.bedrock; non-fatal retrieval error handling with CancelledError re-raise.

Testing

• 68 unit tests pass (KB tool/provider + serializer coalescing + existing).
• Verified end-to-end against a live managed KB in us-west-2:
  • agent.run() with BedrockChatClient + the provider — single valid conversation, correct grounded response.
  • Agentic retrieval returns real source URLs and no fabricated scores.
  • generateResponse=False confirmed to suppress the streamed generation events.
• SDK requirement: boto3 >= 1.43.32
• Microsoft CLA: previously agreed (Amazon).

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This is not a breaking change

[Vidyadhar Pogul (PVidyadhar)]

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 user roles with BedrockChatClient: the provider now injects retrieved context via context.extend_instructions(), so it lands in Bedrock's system field (the prompts list in _prepare_bedrock_messages) rather than the conversation array. This avoids producing consecutive user turns for any Bedrock model. Verified end-to-end with agent.run() + BedrockChatClient against a live KB.

cc Evan Mattson (@moonbox3) Eduard van Valkenburg (@eavanvalkenburg) — thanks for the earlier reviews on #7066.

[Copilot]

🟡 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.

[Vidyadhar Pogul (PVidyadhar)]

Thanks Copilot — addressed all three in 9f8ea75:

1. Prompt-injection / system elevation — Good catch, and it aligns with Evan Mattson (@moonbox3)'s original suggestion. Reverted to keeping retrieved passages as an untrusted user-role message (consistent with azure-cosmos-memory's convention of not elevating retrieved content to instructions). Solved Bedrock's role-alternation requirement by coalescing adjacent user-role messages in _prepare_bedrock_messages — assistant turns are intentionally left unmerged to preserve tool-use/tool-result pairing. Added unit tests for both the coalescing and the no-coalesce-across-assistant cases, and verified end-to-end via agent.run() with BedrockChatClient against a live KB.

2. Stale uv.lock — Regenerated python/uv.lock; the agent-framework-bedrock metadata now records boto3/botocore >=1.43.32.

3. AGENTS.md — Added BedrockKnowledgeBaseTool and BedrockKnowledgeBaseProvider to the package's Main Classes list.

67 unit tests pass; lint clean.

[Copilot]

🟡 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 generatedResponse and responseEvent and 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 assert context.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

[Vidyadhar Pogul (PVidyadhar)]

Thanks Copilot — both addressed in dfd0d2f:

1. Agentic result schema — Confirmed against the live API: AgenticRetrieveStream results expose content/metadata/sourceRetriever and do not carry score or location, so the old code fabricated score: 0 and blank sources for every agentic result. Now the source URI is read from metadata._source_uri and the score is omitted for agentic results (managed reranking doesn't expose one) — the formatter only renders a score when present. Updated the agentic test mock to the real SDK schema. Verified E2E: agentic results now show real docs.aws.amazon.com source URLs and no fabricated scores.

2. Serializer scope — Good catch. Since tool-role messages map to Bedrock user, the coalescing could have merged tool-result blocks into a preceding user text turn. Restricted coalescing to messages whose original role is user, so tool-result turns are never merged — function-call/tool-result serialization is unchanged. Added a regression test covering the user → assistant(toolUse) → tool(toolResult) path to confirm the tool-result stays a distinct turn.

68 unit tests pass; lint clean.

[Copilot]

🟡 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:AgenticRetrieveStream does 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 with Resource: "*", while retaining the KB ARN for bedrock:Retrieve.
  python/packages/bedrock/samples/README.md:35
• bedrock:AgenticRetrieveStream is not resource-scopable, so granting it only on a knowledge-base ARN produces an implicit deny. Split it into a separate statement with Resource: "*"; keep bedrock:Retrieve scoped to the KB ARN.

• Files reviewed: 13/14 changed files
• Comments generated: 2
• Review effort level: Balanced

[Vidyadhar Pogul (PVidyadhar)]

Thanks Copilot — both addressed:

1. Unnecessary response generation (8da7c32) — Confirmed against the live API: omitting generateResponse causes AgenticRetrieveStream to generate a full answer (331 streamed responseEvents), which this tool discards since it only formats passages. Set generateResponse=False explicitly (verified: 0 responseEvents), avoiding the wasted generation latency/cost. Added a test asserting the flag is sent.

2. Stale PR description — Good catch, updated the description to match the actual implementation: the role-alternation fix keeps retrieved content as an untrusted user-role message and solves alternation by coalescing adjacent user turns in _prepare_bedrock_messages (scoped so tool-result turns are never merged). The earlier extend_instructions rationale was from an intermediate revision and no longer applies.

68 unit tests pass; lint clean.

[Copilot]

🟡 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__.py and its .pyi currently omit both names, so from agent_framework.amazon import BedrockKnowledgeBaseTool fails 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 Retrieve location 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:AgenticRetrieveStream has 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. Keep bedrock:Retrieve scoped to the KB ARN, but grant AgenticRetrieveStream in a separate statement with Resource: "*".
  python/packages/bedrock/samples/README.md:37
• bedrock:AgenticRetrieveStream is not resource-scopable, so this sample policy denies that action and the default tool never performs agentic retrieval. Split it into a statement with Resource: "*", while retaining the Knowledge Base ARN restriction for bedrock: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-67 requires every provider-serialization change to update docs/specs/004-python-function-calling-loop.md and 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

[Eduard van Valkenburg (eavanvalkenburg)]

please use the default PR template Vidyadhar Pogul (@PVidyadhar)