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

[Vidyadhar Pogul (PVidyadhar)]

feat: add Amazon Bedrock Knowledge Base tool and context provider

Motivation & Context

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

Description & Review Guide

• What are the major changes?

  • BedrockKnowledgeBaseTool — subclasses FunctionTool with agentic retrieval (AgenticRetrieveStream) and automatic fallback to standard Retrieve. Can be passed directly to any Agent or ChatClient.
  • BedrockKnowledgeBaseProvider — subclasses ContextProvider with before_run() for automatic context injection using standard managed retrieval on every agent invocation.
  • Samples demonstrating both patterns with Agent Framework primitives.

• What is the impact of these changes?

  • Additive only — new files in python/packages/bedrock/, no existing code modified
  • Users get two integration patterns: explicit tool calling (with agentic retrieval) or automatic context injection (standard retrieval)

• What do you want reviewers to focus on?

  • Whether the FunctionTool subclass and ContextProvider subclass match the framework's conventions
  • The before_run() integration using extend_messages()

Related Issue

N/A — new feature adding Amazon Bedrock Knowledge Base integration.

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.

Testing

• 49 unit tests pass (14 for KB tool/provider, 35 existing)
• E2E tested against live managed KB in us-west-2
• tool.invoke() returns formatted retrieval results with scores and sources
• provider.before_run() injects context via extend_messages() into SessionContext
• SDK requirement: boto3 >= 1.43.32
• Microsoft CLA: pre-approved (cla.opensource.microsoft.com)

[Copilot]

Pull request overview

Adds Amazon Bedrock Knowledge Base retrieval capabilities to the Python tools package, enabling both explicit tool-based retrieval and automatic context injection via a ContextProvider.

Changes:

• Introduces BedrockKnowledgeBaseTool with async run() + get_tool_definition() for on-demand KB retrieval.
• Introduces BedrockKnowledgeBaseProvider (ContextProvider.before_run) to inject KB context automatically per agent invocation.
• Adds a design/usage doc plus unit tests for the new tool/provider.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 17 comments.

File	Description
python/packages/tools/agent_framework_tools/bedrock_knowledge_base.py	Adds the Bedrock Knowledge Base tool implementation and tool definition metadata.
python/packages/tools/agent_framework_tools/bedrock_knowledge_base_provider.py	Adds a context provider that retrieves and injects KB passages before runs.
python/packages/tools/tests/test_bedrock_knowledge_base.py	Adds unit tests covering tool retrieval, default config behavior, and provider retrieval.
python/packages/tools/agent_framework_tools/BEDROCK_MANAGED_KB.md	Adds documentation describing usage, configuration, and permissions.

[Vidyadhar Pogul (PVidyadhar)]

@microsoft-github-policy-service agree company="Amazon"

[Eduard van Valkenburg (eavanvalkenburg)]

This is in the wrong place, please move this into the Bedrock package instead. I will look at the implementation then

[Copilot]

🟡 Changes recommended

The provider currently injects retrieved KB content as a system message and lacks non-fatal error handling for retrieval failures, creating avoidable prompt-injection and reliability risks.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

Suppressed comments (4)

python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py:121

• Injected knowledge-base passages are untrusted input and should not be added as a system message (higher-priority prompt-injection risk). Other context providers inject retrieved context as user messages.

        # Inject as a system message via extend_messages
        context_message = Message(role="system", contents=[f"{self.context_prompt}\n\n{retrieved_context}"])
        context.extend_messages(self, [context_message])

python/packages/bedrock/BEDROCK_MANAGED_KB.md:22

• The usage example claims tool.invoke() returns a list of Content items, but BedrockKnowledgeBaseTool formats results into a single string. Update the comment so the docs match the actual return type.

result = asyncio.run(tool.invoke(arguments={"query": "What are the compliance requirements?"}))
print(result)  # List of Content items with retrieval results

python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py:131

• A boto3 failure during retrieval will currently raise and fail the entire agent invocation. Context providers in this repo typically treat retrieval failures as non-fatal and skip injection instead.

        response = await asyncio.to_thread(
            lambda: self._client.retrieve(
                knowledgeBaseId=self.knowledge_base_id,
                retrievalQuery={"text": query},
                retrievalConfiguration={"managedSearchConfiguration": {"numberOfResults": self.number_of_results}},

python/packages/bedrock/tests/test_bedrock_knowledge_base.py:225

• Add an assertion that the injected context message uses the user role, to prevent regressions back to system (prompt-injection risk) and to align with other context providers' behavior.

        assert len(injected) == 1
        assert "Relevant passage" in injected[0].text
        assert "s3://b/doc" in injected[0].text

• Files reviewed: 10/10 changed files
• Comments generated: 1
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Copilot]

🟡 Changes recommended

There are a few correctness and reliability issues (client handling and cancellation swallowing) plus missing end-to-end FunctionTool.invoke() coverage that should be addressed before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

Suppressed comments (5)

python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py:118

• before_run swallows all exceptions from retrieval; this will also swallow task cancellation (asyncio.CancelledError) and can prevent timely shutdown. Re-raise CancelledError explicitly before handling other exceptions.

        try:
            retrieved_context = await self._retrieve(input_text)
        except Exception:
            return

python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:128

• The agentic fallback catches Exception, which can inadvertently catch task cancellation and continue into the standard retrieval path. Re-raise asyncio.CancelledError so cancellations propagate correctly.

            except Exception as e:
                logger.debug("Agentic retrieval failed, falling back: %s", e)

        results = await asyncio.to_thread(self._standard_retrieve, query)
        return self._format_results(results)

python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:97

• client is checked via truthiness, so a valid client object that defines __bool__/__len__ as falsey would be ignored and replaced by a new boto3 client. Prefer an explicit is not None check so the provided client is always respected.

This issue also appears on line 124 of the same file.

        if client:
            self._client = client
        else:
            self._client = boto3.client(

python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py:81

• client is checked via truthiness, so a valid client object that defines __bool__/__len__ as falsey would be ignored and replaced by a new boto3 client. Prefer an explicit is not None check so the provided client is always respected.

This issue also appears on line 115 of the same file.

        if client:
            self._client = client
        else:
            self._client = boto3.client(

python/packages/bedrock/tests/test_bedrock_knowledge_base.py:50

• These tests validate the underlying _retrieve() helper, but they don't exercise the public FunctionTool.invoke() path (argument validation + Content parsing) that users will call. Updating at least one test to call tool.invoke(arguments={...}) would cover the end-to-end tool integration.

        result = asyncio.run(tool._retrieve(query="test query"))
        assert "Result 1" in result
        assert "Result 2" in result
        assert "s3://b/k" in result
        assert "0.950" in result

• Files reviewed: 10/10 changed files
• Comments generated: 0 new
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Copilot]

🟡 Changes recommended

Several docs/samples use an incorrect Agent API shape (won’t run as written), and the provider’s error handling/imports should be adjusted for diagnosability and public-API stability.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

Suppressed comments (9)

python/packages/bedrock/tests/test_bedrock_knowledge_base.py:252

• Prefer importing SessionContext from the public agent_framework API rather than the internal agent_framework._sessions module, to avoid coupling tests to non-public internals.

        from agent_framework._sessions import SessionContext
        from agent_framework_bedrock._knowledge_base_provider import BedrockKnowledgeBaseProvider

python/packages/bedrock/tests/test_bedrock_knowledge_base.py:212

• Prefer importing SessionContext from the public agent_framework API rather than the internal agent_framework._sessions module, to avoid coupling tests to non-public internals.

        from agent_framework import Message
        from agent_framework._sessions import SessionContext

python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py:12

• BedrockKnowledgeBaseProvider is importing ContextProvider, AgentSession, and SessionContext from the internal agent_framework._sessions module. These types are part of the public agent_framework API and other packages in this repo import them from there; using the internal module increases coupling to non-public internals.

from agent_framework import Message
from agent_framework._sessions import AgentSession, ContextProvider, SessionContext
from agent_framework._telemetry import get_user_agent

python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py:120

• The provider swallows all exceptions during retrieval without any logging, which makes configuration/permission failures very hard to diagnose in production while still returning empty context.

        except Exception:
            return

python/packages/bedrock/samples/bedrock_kb_tool.py:48

• This sample uses the old Agent API (chat_client=... and agent.invoke(..., input_message=...)). In the current framework, Agent expects client=... and you should call await agent.run(<message>, session=...).

    agent = Agent(
        name="KnowledgeAssistant",
        instructions="You are a helpful assistant. Use the knowledge base tool to answer questions about the company.",
        chat_client=chat_client,
        tools=[kb_tool],  # FunctionTool subclass, works with any ChatClient

python/packages/bedrock/samples/bedrock_kb_context_provider.py:48

• This sample uses the old Agent API (chat_client=... and agent.invoke(..., input_message=...)). In the current framework, Agent expects client=... and you should call await agent.run(<message>, session=...).

    agent = Agent(
        name="ContextualAssistant",
        instructions="You are a helpful assistant that answers based on provided context.",
        chat_client=chat_client,
        context_providers=[kb_provider],  # ContextProvider subclass, injects context on every run

python/packages/bedrock/BEDROCK_MANAGED_KB.md:23

• The README example constructs Agent(tools=[tool]), but agent_framework.Agent requires a client argument; this snippet as-is will raise at runtime.

# As a FunctionTool, pass directly to an Agent:
agent = Agent(tools=[tool])

python/packages/bedrock/tests/test_bedrock_knowledge_base.py:9

• Tests should prefer importing public API types from agent_framework (e.g., ContextProvider) instead of the internal agent_framework._sessions module, to avoid coupling to non-public internals.

This issue also appears in the following locations of the same file:

• line 211
• line 251

from agent_framework import FunctionTool
from agent_framework._sessions import ContextProvider

python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:62

• The usage example in the BedrockKnowledgeBaseTool docstring constructs Agent(tools=[tool]), but agent_framework.Agent requires a client argument; as written this example will raise TypeError and is misleading.

    Usage:
        from agent_framework_bedrock import BedrockKnowledgeBaseTool

        tool = BedrockKnowledgeBaseTool(knowledge_base_id="YOUR_KB_ID")
        agent = Agent(tools=[tool])

• Files reviewed: 10/10 changed files
• Comments generated: 0 new
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Copilot]

🟡 Changes recommended

A few package-consistency/operational gaps remain in the new Bedrock KB modules (feature-usage telemetry, logger naming consistency, and an intra-package import convention).

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

Suppressed comments (7)

python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:120

• This tool performs Bedrock API calls but does not currently record Bedrock feature usage (mark_feature_used(FeatureIndex.BEDROCK)), unlike the existing Bedrock chat/embedding clients; this reduces observability for feature adoption and diagnostics.

        if self.use_agentic_retrieval:

python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py:117

• This provider calls Bedrock Retrieve but does not record Bedrock feature usage (mark_feature_used(FeatureIndex.BEDROCK)), unlike the existing Bedrock chat/embedding clients; this reduces observability for feature adoption and diagnostics.

        # Retrieve from knowledge base (non-fatal — agent continues without context on failure)
        try:

python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:13

• To keep Bedrock feature-usage telemetry consistent with the rest of the bedrock package (e.g., _chat_client.py and _embedding_client.py), this module should import mark_feature_used and FeatureIndex so the tool can emit FeatureIndex.BEDROCK when it makes Bedrock API calls.

This issue also appears on line 120 of the same file.

from agent_framework import FunctionTool
from agent_framework._telemetry import get_user_agent
from pydantic import BaseModel, Field

python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:27

• Other modules in this package log under the stable "agent_framework.bedrock" logger name (e.g., agent_framework_bedrock/_chat_client.py and _embedding_client.py). Using name here makes log routing/filtering inconsistent across the bedrock package.

logger = logging.getLogger(__name__)

python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py:13

• To keep Bedrock feature-usage telemetry consistent with the rest of the bedrock package, import mark_feature_used and FeatureIndex so the provider can record FeatureIndex.BEDROCK when it retrieves from Bedrock.

This issue also appears on line 116 of the same file.

from agent_framework import AgentSession, ContextProvider, Message, SessionContext
from agent_framework._telemetry import get_user_agent

python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py:27

• This intra-package import should use a relative import (consistent with other agent_framework_bedrock modules, e.g., "from ._feature_usage import FeatureIndex") to avoid issues with refactors and to match established package conventions.

from agent_framework_bedrock._knowledge_base import _get_source_uri

python/packages/bedrock/agent_framework_bedrock/_knowledge_base_provider.py:29

• Other modules in this package log under the stable "agent_framework.bedrock" logger name (e.g., agent_framework_bedrock/_chat_client.py and _embedding_client.py). Using name here makes log routing/filtering inconsistent across the bedrock package.

logger = logging.getLogger(__name__)

• Files reviewed: 10/10 changed files
• Comments generated: 0 new
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Copilot]

🟡 Changes recommended

The new tool can double-count telemetry on agentic retrieval fallback paths, which should be fixed before merge.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

Suppressed comments (1)

python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py:135

• mark_feature_used(FeatureIndex.BEDROCK) is called twice when use_agentic_retrieval=True and the agentic path returns no results or throws (it’s called once before the agentic attempt and again before the standard fallback). This can double-count telemetry for a single tool invocation in common fallback scenarios.

        mark_feature_used(FeatureIndex.BEDROCK)
        results = await asyncio.to_thread(self._standard_retrieve, query)

• Files reviewed: 10/10 changed files
• Comments generated: 0 new
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Copilot]

🟡 Changes recommended

The IAM policy JSON examples in the new documentation are not valid standalone IAM policy documents and should be corrected to avoid users copying broken configurations.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

Suppressed comments (2)

python/packages/bedrock/samples/README.md:33

• The IAM example under "Required IAM Permissions" is not a complete IAM policy document (it’s missing the standard Version + Statement wrapper), so users copying it into IAM will get an invalid policy JSON. Consider providing a full policy document or explicitly labeling this as a single statement snippet.

```json
{
    "Effect": "Allow",
    "Action": [
        "bedrock:Retrieve",

python/packages/bedrock/BEDROCK_MANAGED_KB.md:57

• The IAM example JSON here is not a complete IAM policy document (it’s missing the standard Version + Statement wrapper). As written, it can’t be pasted directly into IAM as a policy without modification.

```json
{
  "Effect": "Allow",
  "Action": [
    "bedrock:Retrieve",

• Files reviewed: 10/10 changed files
• Comments generated: 0 new
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Copilot]

🟢 Ready to approve

The changes are additive, follow existing Agent Framework patterns for tools/providers, and include focused unit tests and documentation.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

• Files reviewed: 10/10 changed files
• Comments generated: 0 new
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Vidyadhar Pogul (PVidyadhar)]

Requesting re-review Evan Mattson (@moonbox3) , Eduard van Valkenburg (@eavanvalkenburg) . thanks for your comments!