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

[Vidyadhar Pogul (PVidyadhar)]

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

• What are the major changes?
  • BedrockKnowledgeBaseTool — subclasses FunctionTool with agentic retrieval (AgenticRetrieveStream; query decomposition + managed reranking) and automatic fallback to standard Retrieve. Passable directly to any Agent or ChatClient. generateResponse=False is set so the tool returns passages only and the agent's own model generates the answer.
  • BedrockKnowledgeBaseProvider — subclasses ContextProvider; its before_run() retrieves passages and injects them as an untrusted user-role message (same convention as the azure-cosmos-memory provider, avoiding elevation of retrieved content to system instructions).
  • Both classes exported from the public agent_framework.amazon namespace.
• What is the impact of these changes?
  • Additive: new files under python/packages/bedrock/ plus namespace exports. No change to shared serialization or the function-calling loop.
• What do you want reviewers to focus on?
  • The FunctionTool / ContextProvider subclassing conventions and the untrusted user-role injection in before_run().

Related Issue

N/A — new feature. No existing open issue or PR (supersedes closed #7066).

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 PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
• 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]

🔵 Needs a closer look

Provider retrieval failures need a production-visible operational signal to prevent silent loss of grounding.

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_provider.py:126

• A credentials, region, throttling, or service failure is only visible at DEBUG level while the agent continues without its expected KB grounding. Under normal production logging this looks like a successful provider run and can yield ungrounded answers with no operational signal. Emit a warning, as the other non-fatal retrieval providers do, while still allowing the run to continue.

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

[Copilot]

🟡 Changes recommended

Fail-open retrieval errors are hidden at DEBUG level, making loss of KB grounding operationally invisible.

Get a fresh assessment by requesting another Copilot review.

Review details

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

[Copilot]

🔵 Needs a closer look

Two usage examples call BedrockChatClient with unsupported constructor arguments.

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:17

• This documented example raises TypeError because BedrockChatClient does not accept an options constructor argument; model_id is also not a BedrockChatOptions key. Pass the model through the client's model parameter.
  python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:156
• This usage example cannot run: BedrockChatClient.__init__ has no options parameter, and the chat option key is model, not model_id. Construct the client with its documented model argument instead.

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

[Vidyadhar Pogul (PVidyadhar)]

Addressed the latest Copilot review (736d27e): the two usage snippets in _knowledge_base.py (docstring) and BEDROCK_MANAGED_KB.md used BedrockChatClient(options=BedrockChatOptions(model_id=...)), which raises TypeError — BedrockChatClient takes a model argument (no options param, and the key is model not model_id). Both now use BedrockChatClient(model="..."), matching the client's own docstring and AGENTS.md. ruff + tests still green.

[Copilot]

🔵 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.type also includes AUDIO and VIDEO. Those payloads do not use content.text, so they currently fall through to line 94 and produce blank passages just like the IMAGE case 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

[Vidyadhar Pogul (PVidyadhar)]

Addressed the latest Copilot review (commit for AUDIO/VIDEO): confirmed against the botocore service model that RetrievalResultContent.type is TEXT | IMAGE | AUDIO | VIDEO | ROW. Generalized the binary-media handling from IMAGE-only to IMAGE/AUDIO/VIDEO (_BINARY_MEDIA_CONTENT_TYPES), each rendered as a typed placeholder instead of a blank passage; updated the helper docstring and added AUDIO/VIDEO tests. ruff + tests green (21 passed).

[Copilot]

🔵 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: ignore in 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 with Any/a small protocol and use only rule-specific, line-level # pyright: ignore[...] suppressions where unavoidable; the legacy _chat_client.py suppression 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: ignore in 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.py exception.

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 ignores BEDROCK_REGION, BEDROCK_ACCESS_KEY, BEDROCK_SECRET_KEY, BEDROCK_SESSION_TOKEN, and a reusable boto3 session. Consequently, a normally configured BedrockChatClient() 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. BedrockChatClient and the embedding client resolve BEDROCK_REGION and BEDROCK_* 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

[Vidyadhar Pogul (PVidyadhar)]

Addressed the latest review (60b3361) — both points:

1. # type: ignore in source — removed the file-level suppressions from _knowledge_base.py and _knowledge_base_provider.py. Per CODING_STANDARD.md, source is strict-pyright and should use line-level # pyright: ignore[...], not file-level # type: ignore (and the legacy _chat_client.py exception shouldn't be copied). The untyped bedrock-agent-runtime boundary is now typed explicitly (responses annotated dict[str, Any]), with only two targeted # pyright: ignore[reportUnknownMemberType] on the dynamic retrieve / agentic_retrieve_stream calls.

2. Settings/session bypass — the tool and provider no longer construct a bare boto3.client(..., region_name=...). They now resolve region + credentials through the shared Bedrock settings 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 constructors now expose region_name/access_key/secret_key/session_token/boto3_session/env_file_*, so a KB tool and a configured BedrockChatClient() target the same region/credentials by default instead of diverging.

Verified: ruff check/format clean and 71 tests pass (1 integration skipped). Also ran an end-to-end check against a live managed KB confirming both the tool and provider resolve region from BEDROCK_REGION and return results.

[Copilot]

🟡 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