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

python/packages/tools/agent_framework_tools/BEDROCK_MANAGED_KB.md

@@ -0,0 +1,53 @@
+# 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_tools.bedrock_kb import BedrockKnowledgeBaseTool
+
+tool = BedrockKnowledgeBaseTool(
+    knowledge_base_id="YOUR_KB_ID",
+    region="us-east-1",
+)
+results = tool.invoke({"query": "What are the compliance requirements?"})
+for result in results:
+    print(result["content"], result["score"])
+```
+
+## Configuration
+| Variable | Description | Default |
+|---|---|---|
+| KNOWLEDGE_BASE_ID | Bedrock Knowledge Base ID | None |
+| AWS_REGION | AWS region for the KB | us-east-1 |
+| AWS_PROFILE | AWS credentials profile | None |
+| USE_AGENTIC_RETRIEVAL | Enable agentic retrieval | true |
+| MAX_RESULTS | Maximum retrieval results | 5 |
+
+## Features
+- Managed search (no vector store needed)
+- Agentic retrieval with query decomposition + reranking
+- Automatic fallback to plain Retrieve if agentic fails
+- Multi-source support (S3, Web, Confluence, SharePoint)
+- Compatible with Agent Framework tool interface
+
+## SDK Requirements
+- boto3 >= 1.43
+
+## Required IAM Permissions
+```json
+{
+  "Effect": "Allow",
+  "Action": [
+    "bedrock:Retrieve",
+    "bedrock:AgenticRetrieve"
+  ],
+  "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/tools/agent_framework_tools/bedrock_knowledge_base.py

@@ -0,0 +1,184 @@
+"""Amazon Bedrock Knowledge Base retrieval tool for Microsoft Agent Framework.
+
+Provides document retrieval from Amazon Bedrock Managed Knowledge Bases
+for use as a tool in agent workflows.
+
+Usage:
+    from agent_framework_tools.bedrock_knowledge_base import BedrockKnowledgeBaseTool
+
+    kb_tool = BedrockKnowledgeBaseTool(knowledge_base_id="ABCDEFGHIJ")
+    results = await kb_tool.run(query="What is our deployment process?")
+"""
+
+import os
+import logging
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+def _get_source_uri(result: dict) -> str:
+    """Extract source URI from a retrieval result, handling all location types."""
+    location = result.get('location', {})
+    loc_type = location.get('type', '')
+    if loc_type == 'S3' or 's3Location' in location:
+        return location.get('s3Location', {}).get('uri', '')
+    elif loc_type == 'WEB' or 'webLocation' in location:
+        return location.get('webLocation', {}).get('url', '')
+    elif 'confluenceLocation' in location:
+        return location.get('confluenceLocation', {}).get('url', '')
+    elif 'salesforceLocation' in location:
+        return location.get('salesforceLocation', {}).get('url', '')
+    elif 'sharePointLocation' in location:
+        return location.get('sharePointLocation', {}).get('url', '')
+    elif 'customDocumentLocation' in location:
+        return location.get('customDocumentLocation', {}).get('id', '')
+    # Fallback to metadata._source_uri (for agentic results)
+    return result.get('metadata', {}).get('_source_uri', '')
+
+
+class BedrockKnowledgeBaseTool:
+    """Retrieves documents from an Amazon Bedrock Managed Knowledge Base.
+
+    Args:
+        knowledge_base_id: The KB ID. Falls back to KNOWLEDGE_BASE_ID env var.
+        region_name: AWS region. Falls back to AWS_REGION env var or us-east-1.
+        number_of_results: Max results to return. Defaults to 5.
+        knowledge_base_type: 'MANAGED' (recommended) or 'VECTOR'.
+        use_agentic_retrieval: If True, try AgenticRetrieveStream first with fallback to plain Retrieve.
+    """
+
+    name: str = "bedrock_knowledge_base"
+    description: str = (
+        "Retrieves relevant documents from an Amazon Bedrock Knowledge Base. "
+        "Use this to search internal documentation and knowledge sources."
+    )
+
+    def __init__(
+        self,
+        knowledge_base_id: Optional[str] = None,
+        region_name: Optional[str] = None,
+        number_of_results: int = 5,
+        knowledge_base_type: str = "MANAGED",
+        use_agentic_retrieval: Optional[bool] = None,
+    ):
+        self.knowledge_base_id = knowledge_base_id or os.environ.get("KNOWLEDGE_BASE_ID", "")
+        self.region_name = region_name or os.environ.get("AWS_REGION", "us-east-1")
+        self.number_of_results = number_of_results
+        self.knowledge_base_type = knowledge_base_type
+        self.use_agentic_retrieval = use_agentic_retrieval if use_agentic_retrieval is not None else os.environ.get('USE_AGENTIC_RETRIEVAL', 'true').lower() != 'false'
+        self._client = None
+
+    @property
+    def client(self):
+        if self._client is None:
+            try:
+                import boto3
+            except ImportError:
+                raise ImportError(
+                    "boto3 is required for Bedrock Knowledge Base tool. "
+                    "Install with: pip install boto3>=1.41.0"
+                )
+            self._client = boto3.client(
+                "bedrock-agent-runtime", region_name=self.region_name
+            )
+        return self._client
+
+    async def run(self, query: str, **kwargs) -> list[dict[str, Any]]:
+        """Retrieve relevant documents.
+
+        Args:
+            query: The search query.
+
+        Returns:
+            List of results with content, source, and score.
+        """
+        k = kwargs.get("max_results", self.number_of_results)
+
+        # Try agentic retrieval first
+        if self.use_agentic_retrieval:
+            agentic_results = self._agentic_retrieve(query, k)
+            if agentic_results is not None:
+                return agentic_results
+
+        # Fallback to managed/vector retrieve
+        if self.knowledge_base_type == "MANAGED":
+            retrieval_config = {
+                "managedSearchConfiguration": {"numberOfResults": k}
+            }
+        else:
+            retrieval_config = {
+                "vectorSearchConfiguration": {"numberOfResults": k}
+            }
+
+        try:
+            response = self.client.retrieve(
+                knowledgeBaseId=self.knowledge_base_id,
+                retrievalQuery={"text": query},
+                retrievalConfiguration=retrieval_config,
+            )
+
+            results = []
+            for result in response.get("retrievalResults", []):
+                content = result.get("content", {}).get("text", "")
+                source = _get_source_uri(result)
+                score = result.get("score", 0.0)
+                results.append({
+                    "content": content,
+                    "source": source,
+                    "score": score,
+                })
+            return results
+        except Exception as e:
+            logger.error(f"Error retrieving from Bedrock KB: {e}")
+            return []
+
+    def _agentic_retrieve(self, query: str, top_k: int) -> list[dict[str, Any]] | None:
+        """Try agentic retrieval with streaming. Returns list of results or None on failure."""
+        try:
+            response = self.client.agentic_retrieve_stream(
+                knowledgeBaseId=self.knowledge_base_id,
+                messages=[{"content": {"text": query}, "role": "user"}],
+                retrievers=[{
+                    "configuration": {
+                        "knowledgeBase": {
+                            "knowledgeBaseId": self.knowledge_base_id,
+                            "retrievalOverrides": {"maxNumberOfResults": top_k},
+                        }
+                    }
+                }],
+                agenticRetrieveConfiguration={
+                    "foundationModelType": "MANAGED",
+                    "rerankingModelType": "MANAGED",
+                },
+            )
+            results = []
+            for event in response.get("stream", []):
+                if "result" in event and "results" in event["result"]:
+                    for result in event["result"]["results"]:
+                        results.append({
+                            "content": result.get("content", {}).get("text", ""),
+                            "source": _get_source_uri(result),
+                            "score": result.get("score", 0.0),
+                        })
+            return results
+        except Exception as e:
+            logger.debug(f"Agentic retrieval unavailable, will fall back to managed retrieve: {e}")
+            return None
+
+    def get_tool_definition(self) -> dict[str, Any]:
+        """Return the tool definition for registration with an agent."""
+        return {
+            "name": self.name,
+            "description": self.description,
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": "The search query to find relevant documents.",
+                    }
+                },
+                "required": ["query"],
+            },
+        }