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

python/packages/bedrock/BEDROCK_MANAGED_KB.md

@@ -0,0 +1,67 @@
+# Bedrock Managed Knowledge Base Support
+
+## Overview
+Adds an Agent Framework tool that queries Amazon Bedrock Knowledge Bases for managed retrieval within agent pipelines.
+
+## Usage
+```python
+from agent_framework import Agent
+from agent_framework_bedrock import BedrockKnowledgeBaseTool, BedrockChatClient, BedrockChatOptions
+
+tool = BedrockKnowledgeBaseTool(
+    knowledge_base_id="YOUR_KB_ID",
+    region_name="us-east-1",
+)
+
+# As a FunctionTool, pass directly to an Agent:
+agent = Agent(client=BedrockChatClient(options=BedrockChatOptions(model_id="...")), tools=[tool])
+
+# Or invoke directly for testing:
+import asyncio
+result = asyncio.run(tool.invoke(arguments={"query": "What are the compliance requirements?"}))
+print(result)  # List of Content items with retrieval results
+```
+
+## Configuration
+
+All configuration is via constructor parameters:
+
+| Parameter | Description | Default |
+|---|---|---|
+| `knowledge_base_id` | Bedrock Knowledge Base ID (required) | — |
+| `region_name` | AWS region for the KB | `us-east-1` |
+| `number_of_results` | Maximum retrieval results | `5` |
+| `use_agentic_retrieval` | Enable agentic multi-hop retrieval | `True` |
+| `client` | Pre-configured boto3 client (optional) | Auto-created |
+
+## Features
+- Managed search (no vector store needed)
+- **BedrockKnowledgeBaseTool**: Agentic retrieval with query decomposition + reranking, automatic fallback to standard Retrieve
+- **BedrockKnowledgeBaseProvider**: Standard managed retrieval injected as context before each agent run
+- Multi-source support (S3, Web, Confluence, SharePoint)
+- Compatible with Agent Framework FunctionTool and ContextProvider interfaces
+
+## SDK Requirements
+- boto3 >= 1.43.32
+
+## Required IAM Permissions
+```json
+{
+  "Version": "2012-10-17",
+  "Statement": [
+    {
+      "Effect": "Allow",
+      "Action": [
+        "bedrock:Retrieve",
+        "bedrock:AgenticRetrieveStream"
+      ],
+      "Resource": "arn:aws:bedrock:<region>:<account-id>:knowledge-base/<kb-id>"
+    }
+  ]
+}
+```
+
+## References
+- [Build a Managed Knowledge Base](https://docs.aws.amazon.com/bedrock/latest/userguide/kb-build-managed.html)
+- [Retrieve API](https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-retrieve.html)
+- [Agentic Retrieval](https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-agentic.html)

python/packages/bedrock/agent_framework_bedrock/__init__.py

@@ -4,6 +4,8 @@
 
 from ._chat_client import BedrockChatClient, BedrockChatOptions, BedrockGuardrailConfig, BedrockSettings
 from ._embedding_client import BedrockEmbeddingClient, BedrockEmbeddingOptions, BedrockEmbeddingSettings
+from ._knowledge_base import BedrockKnowledgeBaseTool
+from ._knowledge_base_provider import BedrockKnowledgeBaseProvider
 
 try:
     __version__ = importlib.metadata.version(__name__)
@@ -17,6 +19,8 @@
     "BedrockEmbeddingOptions",
     "BedrockEmbeddingSettings",
     "BedrockGuardrailConfig",
+    "BedrockKnowledgeBaseProvider",
+    "BedrockKnowledgeBaseTool",
     "BedrockSettings",
     "__version__",
 ]

python/packages/bedrock/agent_framework_bedrock/_knowledge_base.py

@@ -0,0 +1,192 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Amazon Bedrock Knowledge Base retrieval tool for Agent Framework."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, Annotated, Any, Optional
+
+from agent_framework import FunctionTool
+from agent_framework._telemetry import get_user_agent, mark_feature_used
+from pydantic import BaseModel, Field
+
+from ._feature_usage import FeatureIndex
+
+if TYPE_CHECKING:
+    from botocore.client import BaseClient
+
+try:
+    import boto3
+    from botocore.config import Config as BotoConfig
+except ImportError as e:
+    raise ImportError(
+        "boto3 is required for BedrockKnowledgeBaseTool. "
+        "Install it with: pip install boto3>=1.43.32"
+    ) from e
+
+logger = logging.getLogger("agent_framework.bedrock")
+
+
+def _get_source_uri(result: dict[str, Any]) -> str:
+    """Extract source URI from a retrieval result."""
+    location = result.get("location", {})
+    if "s3Location" in location:
+        return location["s3Location"].get("uri", "")
+    if "webLocation" in location:
+        return location["webLocation"].get("url", "")
+    if "confluenceLocation" in location:
+        return location["confluenceLocation"].get("url", "")
+    if "sharePointLocation" in location:
+        return location["sharePointLocation"].get("url", "")
+    if "customDocumentLocation" in location:
+        return location["customDocumentLocation"].get("id", "")
+    return ""
+
+
+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.")]
+
+
+class BedrockKnowledgeBaseTool(FunctionTool):
+    """Tool that retrieves documents from Amazon Bedrock Knowledge Bases.
+
+    Subclasses FunctionTool so it can be passed directly to any Agent or ChatClient.
+
+    Usage:
+        from agent_framework_bedrock import BedrockKnowledgeBaseTool, BedrockChatClient, BedrockChatOptions
+        from agent_framework import Agent
+
+        tool = BedrockKnowledgeBaseTool(knowledge_base_id="YOUR_KB_ID")
+        agent = Agent(client=BedrockChatClient(options=BedrockChatOptions(model_id="...")), tools=[tool])
+    """
+
+    def __init__(
+        self,
+        *,
+        knowledge_base_id: str,
+        region_name: str = "us-east-1",
+        number_of_results: int = 5,
+        use_agentic_retrieval: bool = True,
+        client: Optional[BaseClient] = None,
+        name: str = "bedrock_knowledge_base",
+        description: str = (
+            "Retrieves relevant documents from an Amazon Bedrock Knowledge Base. "
+            "Use this to answer questions that require specific knowledge or context."
+        ),
+    ) -> None:
+        """Create a Bedrock Knowledge Base tool.
+
+        Args:
+            knowledge_base_id: The Bedrock Knowledge Base ID.
+            region_name: AWS region name.
+            number_of_results: Maximum number of results to return.
+            use_agentic_retrieval: Use AgenticRetrieveStream for query decomposition + reranking.
+            client: Pre-configured bedrock-agent-runtime client. If not provided, one is created.
+            name: Tool name for model registration.
+            description: Tool description for model context.
+        """
+        self.knowledge_base_id = knowledge_base_id
+        self.region_name = region_name
+        self.number_of_results = number_of_results
+        self.use_agentic_retrieval = use_agentic_retrieval
+
+        if client is not None:
+            self._client = client
+        else:
+            self._client = boto3.client(
+                "bedrock-agent-runtime",
+                region_name=self.region_name,
+                config=BotoConfig(user_agent_extra=f"{get_user_agent()} bedrock-kb"),
+            )
+
+        super().__init__(
+            name=name,
+            description=description,
+            func=self._retrieve,
+            input_model=_BedrockKBQueryInput,
+        )
+
+    async def _retrieve(self, query: str) -> str:
+        """Retrieve documents from the knowledge base.
+
+        Args:
+            query: The search query.
+
+        Returns:
+            Formatted string of retrieval results.
+        """
+        mark_feature_used(FeatureIndex.BEDROCK)
+
+        if self.use_agentic_retrieval:
+            try:
+                results = await asyncio.to_thread(self._agentic_retrieve, query)
+                if results:
+                    return self._format_results(results)
+            except asyncio.CancelledError:
+                raise
+            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)
+
+    def _agentic_retrieve(self, query: str) -> list[dict[str, Any]]:
+        """Use AgenticRetrieveStream for query decomposition + managed reranking."""
+        response = self._client.agentic_retrieve_stream(
+            messages=[{"content": {"text": query}, "role": "user"}],
+            retrievers=[{
+                "configuration": {
+                    "knowledgeBase": {
+                        "knowledgeBaseId": self.knowledge_base_id,
+                        "retrievalOverrides": {"maxNumberOfResults": self.number_of_results},
+                    }
+                }
+            }],
+            agenticRetrieveConfiguration={
+                "foundationModelType": "MANAGED",
+                "rerankingModelType": "MANAGED",
+            },
+        )
+        results = []
+        for event in response.get("stream", []):
+            if "result" in event and "results" in event["result"]:
+                for r in event["result"]["results"]:
+                    results.append({
+                        "content": r.get("content", {}).get("text", ""),
+                        "source": _get_source_uri(r),
+                        "score": r.get("score", 0),
+                    })
+        return results
+
+    def _standard_retrieve(self, query: str) -> list[dict[str, Any]]:
+        """Use standard Retrieve API with managed search configuration."""
+        response = self._client.retrieve(
+            knowledgeBaseId=self.knowledge_base_id,
+            retrievalQuery={"text": query},
+            retrievalConfiguration={"managedSearchConfiguration": {"numberOfResults": self.number_of_results}},
+        )
+        results = []
+        for r in response.get("retrievalResults", []):
+            results.append({
+                "content": r.get("content", {}).get("text", ""),
+                "source": _get_source_uri(r),
+                "score": r.get("score", 0),
+            })
+        return results
+
+    @staticmethod
+    def _format_results(results: list[dict[str, Any]]) -> str:
+        """Format retrieval results as a readable string."""
+        if not results:
+            return "No relevant documents found."
+        parts = []
+        for i, r in enumerate(results, 1):
+            source = r.get("source", "")
+            content = r.get("content", "")
+            score = r.get("score", 0)
+            parts.append(f"[{i}] (score: {score:.3f}) {content}\n    Source: {source}")
+        return "\n\n".join(parts)