From 962c5f9d76f3f153b858cbbde8fc19565bfdb75f Mon Sep 17 00:00:00 2001 From: PVidyadhar Date: Sat, 11 Jul 2026 02:30:02 +0000 Subject: [PATCH] feat: add Amazon Bedrock Knowledge Base tool and context provider - Created BedrockKnowledgeBaseTool with async run() + get_tool_definition() - Created BedrockKnowledgeBaseProvider (ContextProvider subclass) with before_run() - Two integration points: standalone tool + automatic context injection - Supports managed search and agentic retrieval with fallback - Unit tests included - Added BEDROCK_MANAGED_KB.md design doc --- .../BEDROCK_MANAGED_KB.md | 53 +++++ .../bedrock_knowledge_base.py | 184 ++++++++++++++++ .../bedrock_knowledge_base_provider.py | 205 ++++++++++++++++++ .../tests/test_bedrock_knowledge_base.py | 68 ++++++ 4 files changed, 510 insertions(+) create mode 100644 python/packages/tools/agent_framework_tools/BEDROCK_MANAGED_KB.md create mode 100644 python/packages/tools/agent_framework_tools/bedrock_knowledge_base.py create mode 100644 python/packages/tools/agent_framework_tools/bedrock_knowledge_base_provider.py create mode 100644 python/packages/tools/tests/test_bedrock_knowledge_base.py diff --git a/python/packages/tools/agent_framework_tools/BEDROCK_MANAGED_KB.md b/python/packages/tools/agent_framework_tools/BEDROCK_MANAGED_KB.md new file mode 100644 index 00000000000..51b35c9899d --- /dev/null +++ b/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:::knowledge-base/" +} +``` + +## 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) diff --git a/python/packages/tools/agent_framework_tools/bedrock_knowledge_base.py b/python/packages/tools/agent_framework_tools/bedrock_knowledge_base.py new file mode 100644 index 00000000000..9c7188914b7 --- /dev/null +++ b/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"], + }, + } diff --git a/python/packages/tools/agent_framework_tools/bedrock_knowledge_base_provider.py b/python/packages/tools/agent_framework_tools/bedrock_knowledge_base_provider.py new file mode 100644 index 00000000000..ee036677694 --- /dev/null +++ b/python/packages/tools/agent_framework_tools/bedrock_knowledge_base_provider.py @@ -0,0 +1,205 @@ +# Copyright (c) Amazon.com, Inc. All rights reserved. + +"""Amazon Bedrock Knowledge Base Context Provider for Microsoft Agent Framework. + +Provides automatic retrieval-augmented context from a Bedrock Managed Knowledge Base +before each model invocation — similar to FoundryMemoryProvider but backed by AWS. + +Usage: + from agent_framework import Agent, AgentSession + from agent_framework_tools.bedrock_knowledge_base_provider import BedrockKnowledgeBaseProvider + + kb_provider = BedrockKnowledgeBaseProvider( + source_id="bedrock-kb", + knowledge_base_id="ABCDEFGHIJ", + region_name="us-west-2", + ) + + agent = Agent(context_providers=[kb_provider]) +""" + +from __future__ import annotations + +import logging +import os +from typing import TYPE_CHECKING, Any, Optional + +from agent_framework import ( + AgentSession, + ContextProvider, + Message, + SessionContext, +) + +if TYPE_CHECKING: + from agent_framework import SupportsAgentRun + +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 BedrockKnowledgeBaseProvider(ContextProvider): + """Context provider that retrieves relevant documents from a Bedrock Knowledge Base. + + Automatically queries the KB with the user's latest message before each model + invocation, injecting retrieved context as system messages. + + Args: + source_id: Unique identifier for this provider instance. + 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 retrieve per query. Defaults to 5. + knowledge_base_type: 'MANAGED' (recommended) or 'VECTOR'. + min_score: Minimum relevance score to include a result. Defaults to 0.0. + """ + + def __init__( + self, + source_id: str = "bedrock-kb", + knowledge_base_id: Optional[str] = None, + region_name: Optional[str] = None, + number_of_results: int = 5, + knowledge_base_type: str = "MANAGED", + min_score: float = 0.0, + ): + super().__init__(source_id=source_id) + 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.min_score = min_score + self._client = None + + @property + def client(self): + if self._client is None: + try: + import boto3 + except ImportError: + raise ImportError( + "boto3 is required for BedrockKnowledgeBaseProvider. " + "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 before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + """Retrieve relevant KB context before model invocation. + + Extracts the latest user message, queries the knowledge base, + and injects retrieved passages as context messages. + """ + if not self.knowledge_base_id: + logger.warning("No knowledge_base_id configured. Skipping KB context.") + return + + # Extract the latest user query from the session + query = self._extract_latest_query(session) + if not query: + return + + # Retrieve from KB + passages = await self._retrieve(query) + if not passages: + return + + # Inject as context + context_text = self._format_context(passages) + context.add_instructions( + f"[Knowledge Base Context from {self.source_id}]\n{context_text}" + ) + + def _extract_latest_query(self, session: AgentSession) -> str: + """Extract the most recent user message as the retrieval query.""" + messages = session.messages if hasattr(session, "messages") else [] + for message in reversed(messages): + if hasattr(message, "role") and message.role == "user": + if hasattr(message, "content"): + content = message.content + if isinstance(content, str): + return content + elif isinstance(content, list): + # Extract text from content parts + for part in content: + if isinstance(part, dict) and part.get("type") == "text": + return part.get("text", "") + elif isinstance(part, str): + return part + return "" + + async def _retrieve(self, query: str) -> list[dict[str, Any]]: + """Query the Bedrock Knowledge Base.""" + if self.knowledge_base_type == "MANAGED": + retrieval_config: dict[str, Any] = { + "managedSearchConfiguration": { + "numberOfResults": self.number_of_results + } + } + else: + retrieval_config = { + "vectorSearchConfiguration": { + "numberOfResults": self.number_of_results + } + } + + try: + response = self.client.retrieve( + knowledgeBaseId=self.knowledge_base_id, + retrievalQuery={"text": query}, + retrievalConfiguration=retrieval_config, + ) + + results = [] + for result in response.get("retrievalResults", []): + score = result.get("score", 0.0) + if score < self.min_score: + continue + content = result.get("content", {}).get("text", "") + source = _get_source_uri(result) + if content: + 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 _format_context(self, passages: list[dict[str, Any]]) -> str: + """Format retrieved passages into a context string.""" + formatted = [] + for i, passage in enumerate(passages, 1): + source = passage.get("source", "unknown") + content = passage.get("content", "") + formatted.append(f"[{i}] {content}\n Source: {source}") + return "\n\n".join(formatted) diff --git a/python/packages/tools/tests/test_bedrock_knowledge_base.py b/python/packages/tools/tests/test_bedrock_knowledge_base.py new file mode 100644 index 00000000000..691be0832c7 --- /dev/null +++ b/python/packages/tools/tests/test_bedrock_knowledge_base.py @@ -0,0 +1,68 @@ +"""Tests for Bedrock Knowledge Base tool and context provider.""" +from unittest.mock import MagicMock, patch +import pytest + + +class TestBedrockKnowledgeBaseTool: + @patch("boto3.client") + def test_run_returns_results(self, mock_boto3): + from agent_framework_tools.bedrock_knowledge_base import BedrockKnowledgeBaseTool + mock_client = MagicMock() + mock_client.retrieve.return_value = { + "retrievalResults": [ + {"content": {"text": "Doc"}, "location": {"s3Location": {"uri": "s3://b/d"}}, "score": 0.9}, + ] + } + mock_boto3.return_value = mock_client + tool = BedrockKnowledgeBaseTool(knowledge_base_id="TEST123") + import asyncio + results = asyncio.run(tool.run(query="test")) + assert len(results) == 1 + assert results[0]["content"] == "Doc" + + @patch("boto3.client") + def test_managed_config_default(self, mock_boto3): + from agent_framework_tools.bedrock_knowledge_base import BedrockKnowledgeBaseTool + mock_client = MagicMock() + mock_client.retrieve.return_value = {"retrievalResults": []} + mock_boto3.return_value = mock_client + tool = BedrockKnowledgeBaseTool(knowledge_base_id="TEST123") + import asyncio + asyncio.run(tool.run(query="test")) + call_kwargs = mock_client.retrieve.call_args.kwargs + assert "managedSearchConfiguration" in call_kwargs["retrievalConfiguration"] + + def test_get_tool_definition(self): + from agent_framework_tools.bedrock_knowledge_base import BedrockKnowledgeBaseTool + tool = BedrockKnowledgeBaseTool(knowledge_base_id="TEST123") + defn = tool.get_tool_definition() + assert defn["name"] == "bedrock_knowledge_base" + assert "query" in defn["parameters"]["properties"] + + +class TestBedrockKnowledgeBaseProvider: + def test_init(self): + from agent_framework_tools.bedrock_knowledge_base_provider import BedrockKnowledgeBaseProvider + provider = BedrockKnowledgeBaseProvider( + source_id="test-kb", + knowledge_base_id="TEST123", + region_name="us-west-2", + ) + assert provider.source_id == "test-kb" + assert provider.knowledge_base_id == "TEST123" + + @patch("boto3.client") + def test_retrieve_returns_passages(self, mock_boto3): + from agent_framework_tools.bedrock_knowledge_base_provider import BedrockKnowledgeBaseProvider + mock_client = MagicMock() + mock_client.retrieve.return_value = { + "retrievalResults": [ + {"content": {"text": "Context doc"}, "location": {"s3Location": {"uri": "s3://b/c"}}, "score": 0.9}, + ] + } + mock_boto3.return_value = mock_client + provider = BedrockKnowledgeBaseProvider(knowledge_base_id="TEST123") + import asyncio + passages = asyncio.run(provider._retrieve("test query")) + assert len(passages) == 1 + assert passages[0]["content"] == "Context doc"