From 97eeb77f2bede007b7d272654dfadeac87451728 Mon Sep 17 00:00:00 2001 From: PVidyadhar Date: Sat, 11 Jul 2026 02:27:15 +0000 Subject: [PATCH] feat: add Amazon Bedrock Knowledge Base tool - Created BedrockKBTool(BaseTool) subclass with _get_declaration() and async run_async() - Supports managed and vector knowledge base types - Agentic retrieval with fallback to standard Retrieve API - Unit tests included - Added BEDROCK_MANAGED_KB.md design doc --- PR_DESCRIPTION.md | 46 ++++ src/google/adk/tools/BEDROCK_MANAGED_KB.md | 55 +++++ src/google/adk/tools/bedrock_kb_tool.py | 207 ++++++++++++++++++ tests/test_bedrock_kb_tool.py | 231 +++++++++++++++++++++ 4 files changed, 539 insertions(+) create mode 100644 PR_DESCRIPTION.md create mode 100644 src/google/adk/tools/BEDROCK_MANAGED_KB.md create mode 100644 src/google/adk/tools/bedrock_kb_tool.py create mode 100644 tests/test_bedrock_kb_tool.py diff --git a/PR_DESCRIPTION.md b/PR_DESCRIPTION.md new file mode 100644 index 0000000000..6ee8143cf0 --- /dev/null +++ b/PR_DESCRIPTION.md @@ -0,0 +1,46 @@ +# feat: add Amazon Bedrock Knowledge Base tool + +**Please ensure you have read the [contribution guide](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) before creating a pull request.** + +### Link to Issue or Description of Change + +**1. Link to an existing issue (if applicable):** + +N/A — new feature. + +**2. Or, if no issue exists, describe the change:** + +**Problem:** +ADK agents currently have no built-in way to retrieve context from Amazon Bedrock Knowledge Bases, requiring users to write custom tool implementations for RAG workflows. + +**Solution:** +Created `BedrockKBTool(BaseTool)` subclass that integrates natively with the ADK agent framework. The tool uses `_get_declaration()` for schema registration and `async run_async()` for execution, matching the standard ADK tool pattern. Supports both managed and vector KB types with agentic retrieval (query decomposition + managed reranking) and automatic fallback. + +### Testing Plan + +**Unit Tests:** + +- [x] I have added or updated unit tests for my change. +- [x] All unit tests pass locally. + +`pytest tests/test_bedrock_kb_tool.py` — all tests pass with mocked boto3 calls. + +**Manual End-to-End (E2E) Tests:** + +Called `tool.run_async(args={'query': '...'})` against a managed KB in us-west-2. Returned 3 results with content and scores. Called exactly as ADK agent framework would invoke it. + +### Checklist + +- [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document. +- [x] I have performed a self-review of my own code. +- [x] I have commented my code, particularly in hard-to-understand areas. +- [x] I have added tests that prove my fix is effective or that my feature works. +- [x] New and existing unit tests pass locally with my changes. +- [x] I have manually tested my changes end-to-end. +- [ ] Any dependent changes have been merged and published in downstream modules. + +### Additional context + +- Required IAM permissions: `bedrock:Retrieve`, `bedrock:AgenticRetrieve` +- SDK requirement: `boto3 >= 1.43` +- Google CLA: Approved (P469013198) diff --git a/src/google/adk/tools/BEDROCK_MANAGED_KB.md b/src/google/adk/tools/BEDROCK_MANAGED_KB.md new file mode 100644 index 0000000000..5d95f9bf31 --- /dev/null +++ b/src/google/adk/tools/BEDROCK_MANAGED_KB.md @@ -0,0 +1,55 @@ +# Bedrock Managed Knowledge Base Support + +## Overview +Adds a Google ADK tool that queries Amazon Bedrock Knowledge Bases for managed retrieval within ADK agents. + +## Usage +```python +from google.adk import Agent +from google.adk.tools import BedrockKnowledgeBaseTool + +kb_tool = BedrockKnowledgeBaseTool(knowledge_base_id="YOUR_KB_ID") +agent = Agent( + name="research_agent", + model="gemini-2.0-flash", + tools=[kb_tool], + instruction="Use the knowledge base to answer questions.", +) +``` + +## 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 ADK BaseTool interface + +## SDK Requirements +- boto3 >= 1.43 +- google-adk >= 0.1 + +## 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/src/google/adk/tools/bedrock_kb_tool.py b/src/google/adk/tools/bedrock_kb_tool.py new file mode 100644 index 0000000000..65765bde90 --- /dev/null +++ b/src/google/adk/tools/bedrock_kb_tool.py @@ -0,0 +1,207 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Amazon Bedrock Knowledge Base retrieval tool for Google ADK. + +Provides a tool that queries Amazon Bedrock Managed Knowledge Bases +for use in ADK agents. + +Usage: + from google.adk.tools.bedrock_kb_tool import BedrockKBTool + + kb_tool = BedrockKBTool(knowledge_base_id="ABCDEFGHIJ") + agent = Agent(tools=[kb_tool]) +""" + +import os +from typing import Any, Optional + +from google.adk.tools.base_tool import BaseTool +from google.genai import types as genai_types + + +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 BedrockKBTool(BaseTool): + """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: Use AgenticRetrieveStream for complex queries with + query decomposition and managed reranking. Falls back to plain Retrieve + on failure. Defaults to True. + """ + + 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, + ): + super().__init__( + name="bedrock_knowledge_base", + description=( + "Retrieves relevant documents from an Amazon Bedrock Knowledge Base. " + "Use this to search for information in the knowledge base." + ), + ) + 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 + + def _get_client(self): + if self._client is None: + try: + import boto3 + except ImportError: + raise ImportError( + "boto3 is required for BedrockKBTool. " + "Install with: pip install boto3>=1.41.0" + ) + self._client = boto3.client( + "bedrock-agent-runtime", region_name=self.region_name + ) + return self._client + + def _get_declaration(self) -> genai_types.FunctionDeclaration: + """Return the function declaration for this tool.""" + return genai_types.FunctionDeclaration( + name=self.name, + description=self.description, + parameters=genai_types.Schema( + type="OBJECT", + properties={ + "query": genai_types.Schema( + type="STRING", + description="The search query to find relevant documents.", + ) + }, + required=["query"], + ), + ) + + def _managed_retrieve(self, query: str) -> dict[str, Any]: + """Retrieve using plain managed Retrieve API.""" + client = self._get_client() + + if self.knowledge_base_type == "MANAGED": + retrieval_config = { + "managedSearchConfiguration": { + "numberOfResults": self.number_of_results + } + } + else: + retrieval_config = { + "vectorSearchConfiguration": { + "numberOfResults": self.number_of_results + } + } + + response = 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": results} + + def _agentic_retrieve(self, query: str) -> dict[str, Any]: + """Retrieve using AgenticRetrieveStream with fallback to plain Retrieve.""" + try: + client = self._get_client() + response = 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": self.number_of_results + }, + } + } + }], + agenticRetrieveConfiguration={ + "foundationModelType": "MANAGED", + "rerankingModelType": "MANAGED", + }, + ) + # Process streaming response + results = [] + for event in response.get("stream", []): + if "result" in event and "results" in event["result"]: + for result in event["result"]["results"]: + 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": results} + except Exception: + # Fall back to plain managed retrieve + return self._managed_retrieve(query) + + async def run_async(self, *, args: dict[str, Any], **kwargs) -> dict[str, Any]: + """Execute the retrieval.""" + query = args.get("query", "") + if not query: + return {"error": "No query provided."} + + try: + if self.use_agentic_retrieval: + return self._agentic_retrieve(query) + return self._managed_retrieve(query) + except Exception as e: + return {"error": f"Error retrieving from Bedrock KB: {e}"} diff --git a/tests/test_bedrock_kb_tool.py b/tests/test_bedrock_kb_tool.py new file mode 100644 index 0000000000..2abbfd3c0d --- /dev/null +++ b/tests/test_bedrock_kb_tool.py @@ -0,0 +1,231 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock, patch + +from google.adk.tools.bedrock_kb_tool import BedrockKBTool +from google.genai import types as genai_types +import pytest + + +class TestBedrockKBToolInit: + """Tests for BedrockKBTool initialization.""" + + def test_init_with_explicit_params(self): + tool = BedrockKBTool( + knowledge_base_id="KB123", + region_name="us-west-2", + number_of_results=10, + knowledge_base_type="VECTOR", + ) + + assert tool.knowledge_base_id == "KB123" + assert tool.region_name == "us-west-2" + assert tool.number_of_results == 10 + assert tool.knowledge_base_type == "VECTOR" + assert tool.name == "bedrock_knowledge_base" + + @patch.dict("os.environ", {}, clear=True) + def test_init_defaults(self): + """With no params and no env vars, region defaults to us-east-1.""" + tool = BedrockKBTool(knowledge_base_id="KB456") + + assert tool.knowledge_base_id == "KB456" + assert tool.region_name == "us-east-1" + assert tool.number_of_results == 5 + assert tool.knowledge_base_type == "MANAGED" + + @patch.dict("os.environ", {"KNOWLEDGE_BASE_ID": "ENV_KB_ID", "AWS_REGION": "eu-west-1"}) + def test_init_from_env_vars(self): + tool = BedrockKBTool() + + assert tool.knowledge_base_id == "ENV_KB_ID" + assert tool.region_name == "eu-west-1" + + @patch.dict("os.environ", {"KNOWLEDGE_BASE_ID": "ENV_KB_ID", "AWS_REGION": "eu-west-1"}) + def test_explicit_params_override_env_vars(self): + tool = BedrockKBTool(knowledge_base_id="EXPLICIT_ID", region_name="ap-southeast-1") + + assert tool.knowledge_base_id == "EXPLICIT_ID" + assert tool.region_name == "ap-southeast-1" + + +class TestBedrockKBToolDeclaration: + """Tests for _get_declaration.""" + + def test_get_declaration_returns_correct_schema(self): + tool = BedrockKBTool(knowledge_base_id="KB123") + declaration = tool._get_declaration() + + assert isinstance(declaration, genai_types.FunctionDeclaration) + assert declaration.name == "bedrock_knowledge_base" + assert "knowledge base" in declaration.description.lower() + assert declaration.parameters.type == "OBJECT" + assert "query" in declaration.parameters.properties + assert declaration.parameters.properties["query"].type == "STRING" + assert declaration.parameters.required == ["query"] + + +class TestBedrockKBToolRunAsync: + """Tests for run_async method.""" + + @pytest.fixture + def mock_boto3_client(self): + """Create a mock client and return it for direct injection into tool._client.""" + mock_client = MagicMock() + return mock_client + + def _make_tool(self, mock_client, **kwargs): + """Create a BedrockKBTool with agentic retrieval disabled and inject mock client.""" + kwargs.setdefault("use_agentic_retrieval", False) + tool = BedrockKBTool(**kwargs) + tool._client = mock_client + return tool + + @pytest.mark.asyncio + async def test_run_async_managed_config(self, mock_boto3_client): + tool = self._make_tool( + mock_boto3_client, + knowledge_base_id="KB_MANAGED", + number_of_results=3, + knowledge_base_type="MANAGED", + ) + + mock_boto3_client.retrieve.return_value = { + "retrievalResults": [ + { + "content": {"text": "Document content 1"}, + "location": {"s3Location": {"uri": "s3://bucket/doc1.pdf"}}, + "score": 0.95, + }, + { + "content": {"text": "Document content 2"}, + "location": {"s3Location": {"uri": "s3://bucket/doc2.pdf"}}, + "score": 0.85, + }, + ] + } + + result = await tool.run_async(args={"query": "test query"}) + + mock_boto3_client.retrieve.assert_called_once_with( + knowledgeBaseId="KB_MANAGED", + retrievalQuery={"text": "test query"}, + retrievalConfiguration={ + "managedSearchConfiguration": {"numberOfResults": 3} + }, + ) + + assert "results" in result + assert len(result["results"]) == 2 + assert result["results"][0]["content"] == "Document content 1" + assert result["results"][0]["source"] == "s3://bucket/doc1.pdf" + assert result["results"][0]["score"] == 0.95 + + @pytest.mark.asyncio + async def test_run_async_vector_config(self, mock_boto3_client): + tool = self._make_tool( + mock_boto3_client, + knowledge_base_id="KB_VECTOR", + number_of_results=7, + knowledge_base_type="VECTOR", + ) + + mock_boto3_client.retrieve.return_value = { + "retrievalResults": [ + { + "content": {"text": "Vector result"}, + "location": {"s3Location": {"uri": "s3://bucket/vec.txt"}}, + "score": 0.9, + }, + ] + } + + result = await tool.run_async(args={"query": "vector search"}) + + mock_boto3_client.retrieve.assert_called_once_with( + knowledgeBaseId="KB_VECTOR", + retrievalQuery={"text": "vector search"}, + retrievalConfiguration={ + "vectorSearchConfiguration": {"numberOfResults": 7} + }, + ) + + assert len(result["results"]) == 1 + assert result["results"][0]["content"] == "Vector result" + + @pytest.mark.asyncio + async def test_run_async_empty_results(self, mock_boto3_client): + tool = self._make_tool(mock_boto3_client, knowledge_base_id="KB_EMPTY") + + mock_boto3_client.retrieve.return_value = {"retrievalResults": []} + + result = await tool.run_async(args={"query": "no matches"}) + + assert result == {"results": []} + + @pytest.mark.asyncio + async def test_run_async_empty_query_returns_error(self, mock_boto3_client): + tool = self._make_tool(mock_boto3_client, knowledge_base_id="KB123") + + result = await tool.run_async(args={"query": ""}) + + assert "error" in result + assert "No query provided" in result["error"] + mock_boto3_client.retrieve.assert_not_called() + + @pytest.mark.asyncio + async def test_run_async_missing_query_returns_error(self, mock_boto3_client): + tool = self._make_tool(mock_boto3_client, knowledge_base_id="KB123") + + result = await tool.run_async(args={}) + + assert "error" in result + assert "No query provided" in result["error"] + mock_boto3_client.retrieve.assert_not_called() + + @pytest.mark.asyncio + async def test_run_async_api_error(self, mock_boto3_client): + tool = self._make_tool(mock_boto3_client, knowledge_base_id="KB_ERROR") + + mock_boto3_client.retrieve.side_effect = Exception( + "AccessDeniedException: Not authorized" + ) + + result = await tool.run_async(args={"query": "will fail"}) + + assert "error" in result + assert "Error retrieving from Bedrock KB" in result["error"] + assert "AccessDeniedException" in result["error"] + + @pytest.mark.asyncio + async def test_run_async_partial_result_fields(self, mock_boto3_client): + """Results with missing optional fields should use defaults.""" + tool = self._make_tool(mock_boto3_client, knowledge_base_id="KB_PARTIAL") + + mock_boto3_client.retrieve.return_value = { + "retrievalResults": [ + { + "content": {"text": "Partial doc"}, + # No location or score + }, + ] + } + + result = await tool.run_async(args={"query": "partial"}) + + assert len(result["results"]) == 1 + assert result["results"][0]["content"] == "Partial doc" + assert result["results"][0]["source"] == "" + assert result["results"][0]["score"] == 0.0