diff --git a/src/google/adk_community/tools/bedrock_kb/__init__.py b/src/google/adk_community/tools/bedrock_kb/__init__.py new file mode 100644 index 0000000..846e954 --- /dev/null +++ b/src/google/adk_community/tools/bedrock_kb/__init__.py @@ -0,0 +1,21 @@ +# 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 tool for Google ADK agents.""" + +from google.adk_community.tools.bedrock_kb.bedrock_kb_tool import ( + bedrock_kb_retrieve, +) + +__all__ = ["bedrock_kb_retrieve"] diff --git a/src/google/adk_community/tools/bedrock_kb/bedrock_kb_tool.py b/src/google/adk_community/tools/bedrock_kb/bedrock_kb_tool.py new file mode 100644 index 0000000..6aefe23 --- /dev/null +++ b/src/google/adk_community/tools/bedrock_kb/bedrock_kb_tool.py @@ -0,0 +1,167 @@ +# 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 function tool that queries Amazon Bedrock Managed Knowledge Bases +for retrieval-augmented generation (RAG) in ADK agents. + +Environment Variables: + KNOWLEDGE_BASE_ID: Bedrock Knowledge Base ID (required if not passed as arg) + AWS_REGION: AWS region (default: us-east-1) + USE_AGENTIC_RETRIEVAL: Enable agentic multi-hop retrieval (default: true) + +Dependencies: + pip install boto3>=1.43.2 + +Usage: + from google.adk_community.tools.bedrock_kb import bedrock_kb_retrieve + from google.adk.agents import Agent + + agent = Agent( + model="gemini-2.0-flash", + tools=[bedrock_kb_retrieve], + instruction="Use bedrock_kb_retrieve to answer questions from the knowledge base.", + ) +""" + +import logging +import os +from typing import Optional + +logger = logging.getLogger(__name__) + + +def _get_source_uri(result: dict) -> 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 "salesforceLocation" in location: + return location["salesforceLocation"].get("url", "") + if "sharePointLocation" in location: + return location["sharePointLocation"].get("url", "") + if "customDocumentLocation" in location: + return location["customDocumentLocation"].get("id", "") + return "" + + +def _get_client(): + """Create boto3 bedrock-agent-runtime client with user-agent.""" + import boto3 + from botocore.config import Config + + region = os.environ.get("AWS_REGION", "us-east-1") + return boto3.client( + "bedrock-agent-runtime", + region_name=region, + config=Config(user_agent_extra="google-adk/bedrock-kb"), + ) + + +def _agentic_retrieve(client, knowledge_base_id: str, query: str, max_results: int) -> list[dict]: + """Use AgenticRetrieveStream for multi-hop reasoning retrieval.""" + response = client.agentic_retrieve_stream( + messages=[{"content": {"text": query}, "role": "user"}], + retrievers=[{ + "configuration": { + "knowledgeBase": { + "knowledgeBaseId": knowledge_base_id, + "retrievalOverrides": {"maxNumberOfResults": max_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(client, knowledge_base_id: str, query: str, max_results: int) -> list[dict]: + """Use standard Retrieve API.""" + response = client.retrieve( + knowledgeBaseId=knowledge_base_id, + retrievalQuery={"text": query}, + retrievalConfiguration={"managedSearchConfiguration": {"numberOfResults": max_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 + + +def bedrock_kb_retrieve( + query: str, + knowledge_base_id: Optional[str] = None, + max_results: int = 5, +) -> str: + """Retrieve relevant documents from an Amazon Bedrock Knowledge Base. + + Searches the knowledge base using either agentic retrieval (multi-hop + reasoning with managed reranking) or standard semantic search, and returns + formatted results with content, source, and relevance score. + + Args: + query: The natural language question or search query. + knowledge_base_id: Bedrock Knowledge Base ID. Defaults to KNOWLEDGE_BASE_ID env var. + max_results: Maximum number of results to return. Defaults to 5. + + Returns: + Formatted string with retrieved passages, sources, and scores. + """ + kb_id = knowledge_base_id or os.environ.get("KNOWLEDGE_BASE_ID", "") + if not kb_id: + return "Error: No knowledge_base_id provided. Set KNOWLEDGE_BASE_ID environment variable or pass it as an argument." + + use_agentic = os.environ.get("USE_AGENTIC_RETRIEVAL", "true").lower() != "false" + client = _get_client() + + results = [] + if use_agentic: + try: + results = _agentic_retrieve(client, kb_id, query, max_results) + except Exception as e: + logger.debug("Agentic retrieval failed, falling back to standard: %s", e) + results = _standard_retrieve(client, kb_id, query, max_results) + else: + results = _standard_retrieve(client, kb_id, query, max_results) + + if not results: + return "No relevant documents found in the knowledge base." + + formatted = [] + for i, r in enumerate(results[:max_results], 1): + formatted.append( + f"[{i}] (Score: {r['score']:.2f}) Source: {r['source']}\n{r['content']}" + ) + return "\n\n---\n\n".join(formatted) diff --git a/tests/tools/bedrock_kb/__init__.py b/tests/tools/bedrock_kb/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/tools/bedrock_kb/test_bedrock_kb_tool.py b/tests/tools/bedrock_kb/test_bedrock_kb_tool.py new file mode 100644 index 0000000..508266c --- /dev/null +++ b/tests/tools/bedrock_kb/test_bedrock_kb_tool.py @@ -0,0 +1,127 @@ +# 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. + +"""Tests for Bedrock Knowledge Base tool.""" + +import os +from unittest.mock import MagicMock, patch + +import pytest + + +class TestBedrockKBRetrieve: + """Tests for bedrock_kb_retrieve function.""" + + def test_retrieve_returns_formatted_results(self): + from google.adk_community.tools.bedrock_kb.bedrock_kb_tool import bedrock_kb_retrieve + + mock_client = MagicMock() + mock_client.retrieve.return_value = { + "retrievalResults": [ + { + "content": {"text": "Amazon Bedrock is a managed service."}, + "score": 0.95, + "location": {"s3Location": {"uri": "s3://bucket/doc.pdf"}}, + }, + { + "content": {"text": "Knowledge bases provide RAG capabilities."}, + "score": 0.82, + "location": {"webLocation": {"url": "https://docs.aws.amazon.com/bedrock"}}, + }, + ] + } + + with patch("google.adk_community.tools.bedrock_kb.bedrock_kb_tool._get_client", return_value=mock_client): + with patch.dict(os.environ, {"USE_AGENTIC_RETRIEVAL": "false"}): + result = bedrock_kb_retrieve("What is Bedrock?", knowledge_base_id="TEST_KB") + + assert "Amazon Bedrock is a managed service" in result + assert "Knowledge bases provide RAG capabilities" in result + assert "s3://bucket/doc.pdf" in result + assert "0.95" in result + + def test_agentic_retrieve_with_fallback(self): + from google.adk_community.tools.bedrock_kb.bedrock_kb_tool import bedrock_kb_retrieve + + mock_client = MagicMock() + # Agentic fails + mock_client.agentic_retrieve_stream.side_effect = Exception("Not supported") + # Fallback to standard + mock_client.retrieve.return_value = { + "retrievalResults": [ + {"content": {"text": "Fallback result."}, "score": 0.7, "location": {}}, + ] + } + + with patch("google.adk_community.tools.bedrock_kb.bedrock_kb_tool._get_client", return_value=mock_client): + with patch.dict(os.environ, {"USE_AGENTIC_RETRIEVAL": "true"}): + result = bedrock_kb_retrieve("test query", knowledge_base_id="TEST_KB") + + assert "Fallback result" in result + mock_client.agentic_retrieve_stream.assert_called_once() + mock_client.retrieve.assert_called_once() + + def test_no_kb_id_returns_error(self): + from google.adk_community.tools.bedrock_kb.bedrock_kb_tool import bedrock_kb_retrieve + + with patch.dict(os.environ, {}, clear=True): + result = bedrock_kb_retrieve("test", knowledge_base_id="") + + assert "Error" in result + assert "KNOWLEDGE_BASE_ID" in result + + def test_no_results_returns_message(self): + from google.adk_community.tools.bedrock_kb.bedrock_kb_tool import bedrock_kb_retrieve + + mock_client = MagicMock() + mock_client.retrieve.return_value = {"retrievalResults": []} + + with patch("google.adk_community.tools.bedrock_kb.bedrock_kb_tool._get_client", return_value=mock_client): + with patch.dict(os.environ, {"USE_AGENTIC_RETRIEVAL": "false"}): + result = bedrock_kb_retrieve("unknown topic", knowledge_base_id="TEST_KB") + + assert "No relevant documents found" in result + + def test_agentic_retrieve_success(self): + from google.adk_community.tools.bedrock_kb.bedrock_kb_tool import bedrock_kb_retrieve + + mock_client = MagicMock() + mock_client.agentic_retrieve_stream.return_value = { + "stream": [ + { + "result": { + "results": [ + {"content": {"text": "Agentic result."}, "score": 0.9, "location": {"s3Location": {"uri": "s3://b/k"}}}, + ] + } + } + ] + } + + with patch("google.adk_community.tools.bedrock_kb.bedrock_kb_tool._get_client", return_value=mock_client): + with patch.dict(os.environ, {"USE_AGENTIC_RETRIEVAL": "true"}): + result = bedrock_kb_retrieve("test", knowledge_base_id="TEST_KB") + + assert "Agentic result" in result + mock_client.agentic_retrieve_stream.assert_called_once() + mock_client.retrieve.assert_not_called() + + def test_user_agent_is_set(self): + from google.adk_community.tools.bedrock_kb.bedrock_kb_tool import _get_client + + with patch("boto3.client") as mock_boto: + _get_client() + call_kwargs = mock_boto.call_args + config = call_kwargs.kwargs.get("config") or call_kwargs[1].get("config") + assert "google-adk/bedrock-kb" in config.user_agent_extra