Agent Development Kit (ADK) を利用してAmazon Bedrock KBをtoolsとして利用してみた
はじめに
話題の Google Agent Development Kit(ADK) を使って、AWS Bedrockのナレッジベースとの連携方法を試してみました。
RAGなどの構築でAWSを利用している方(私もその一人です)は、「これまでの資産を活かしつつADKも活用できないかな?」と思ったことがあるのではないでしょうか。そこで今回は、ADKとAWS Bedrock Knowledge Baseの連携にチャレンジしてみた実装例をシェアします。
対象読者
この記事は以下のような方に向けて書いています。
- AWS Bedrockのナレッジベース機能について基本的な知識がある方
- Google ADKを使ったエージェント開発に興味がある方
前提条件
このプロジェクトを実行するには、以下の準備が必要です。
環境要件
- Python 3.13以上
- AWS アカウントとAmazon Bedrockへのアクセス権
- Amazon Bedrock Knowledge Baseの設定済み環境
前提条件
- AWS Bedrockのナレッジベースが設定済みであること(ナレッジベースのIDは後ほど環境変数で指定)
- Pythonのパッケージ管理には
uv
を使用
プロジェクトのセットアップ
プロジェクトの初期化
まずは新しいディレクトリを作成し、プロジェクトを作成します。
mkdir google-adk-bedrock-kb
cd google-adk-bedrock-kb
uvの初期化とADKのインストール
今回はパッケージマネージャとしてuv
を使います。
uv init
uv add google-adk
AWS SDKのboto3もインストールしておきます。
uv add boto3
ディレクトリ構造の作成
次のようなフォルダ構成を作成します。
google-adk-bedrock-kb/
└── agent1/
├── __init__.py
├── agent.py
├── .env
└── tools/
├── __init__.py
└── bedrock_kb.py
以下のコマンドでディレクトリを作成しましょう。
mkdir -p agent1/tools
touch agent1/__init__.py
touch agent1/tools/__init__.py
環境変数の設定
プロジェクトのagent1ディレクトリ内に.envファイルを作成し、以下のように環境変数を設定します。
今回のLLMは、Geminiを利用します。
# .env
GOOGLE_GENAI_USE_VERTEXAI=FALSE
GOOGLE_API_KEY=PASTE_YOUR_ACTUAL_API_KEY_HERE
KNOWLEDGE_BASE_ID=your-knowledge-base-id
AWS_REGION=region
# AWS認証情報(必要に応じて)
AWS_ACCESS_KEY_ID=your-access-key
AWS_SECRET_ACCESS_KEY=your-secret-key
コードの実装
Bedrock KB検索ツールの実装
まず、agent1/tools/bedrock_kb.py
ファイルを作成し、以下のコードを実装します。
import os
from typing import Dict, List, Any
import boto3
from google.adk.tools import FunctionTool
def get_bedrock_client():
"""Bedrockクライアントを初期化して返します。"""
region = os.environ.get("AWS_REGION")
return boto3.client("bedrock-agent-runtime", region_name=region)
def validate_environment():
"""必要な環境変数が設定されているか検証します。"""
knowledge_base_id = os.environ.get("KNOWLEDGE_BASE_ID")
if knowledge_base_id is None:
raise ValueError("Environment variable KNOWLEDGE_BASE_ID is not set")
return knowledge_base_id
def format_retrieval_results(response: Dict[str, Any]) -> List[Dict[str, Any]]:
"""APIレスポンスから検索結果を抽出してフォーマットします。"""
results = []
for result in response.get("retrievalResults", []):
content = result.get("content", {}).get("text", "")
source = result.get("location", {}).get("s3Location", {}).get("uri", "不明")
score = result.get("score", 0)
# メタデータの抽出
metadata = {}
for attr in result.get("metadata", {}).get("attributes", []):
key = attr.get("key", "")
value = attr.get("value", {}).get("text", "")
if key and value:
metadata[key] = value
results.append({
"content": content,
"source": source,
"score": score,
"metadata": metadata,
})
return results
def bedrock_kb_retrieval(query: str) -> Dict[str, Any]:
"""Amazon Bedrock Knowledge Basesから情報を検索します。
ユーザーが特定の情報やドキュメントについて質問した場合にこのツールを使用してください。
ナレッジベースに保存されている情報に基づいて回答を提供します。
Args:
query (str): 検索クエリ。ユーザーの質問や検索したい内容。
Returns:
dict: 検索結果を含む辞書。以下の形式:
{
"status": "success" または "error",
"results": [
{
"content": "取得したコンテンツ",
"source": "情報のソース",
"score": 検索スコア,
"metadata": {属性: 値, ...}
},
...
],
"error_message": エラーメッセージ(エラー発生時のみ)
}
"""
try:
# 環境変数の検証
knowledge_base_id = validate_environment()
# Bedrockクライアントの初期化
bedrock_runtime = get_bedrock_client()
# 検索パラメータの設定
retrieve_params = {
"knowledgeBaseId": knowledge_base_id,
"retrievalQuery": {
"text": query,
},
"retrievalConfiguration": {
"vectorSearchConfiguration": {
"numberOfResults": 7,
},
},
}
# 検索を実行
response = bedrock_runtime.retrieve(**retrieve_params)
# 結果をフォーマット
results = format_retrieval_results(response)
return {
"status": "success",
"results": results,
}
except ValueError as e:
# 環境変数関連のエラー
return {
"status": "error",
"results": [],
"error_message": str(e)
}
except Exception as e:
# その他の例外
error_msg = f"検索中にエラーが発生しました: {str(e)}"
return {
"status": "error",
"results": [],
"error_message": error_msg
}
bedrock_kb_tool = FunctionTool(func=bedrock_kb_retrieval)
このコードでは、AWS Bedrock Knowledge Baseに対するクエリを実行するFunctionToolを定義しています。
エージェントの実装
次に、agent1/agent.py
ファイルを作成し、以下のコードを実装します。
from google.adk.agents import Agent
from .tools.bedrock_kb import bedrock_kb_tool
root_agent = Agent(
model="gemini-2.0-flash-exp",
name="root_agent",
description="This is the root agent for the Bedrock KB project.",
instruction="""質問に答える際には、可能な限り Amazon Bedrock Knowledge Base から情報を検索してください。
**ユーザーが特定のトピックについて質問した場合、'bedrock_kb_retrieval' ツールを使用して関連情報を取得してください。**
**もしツールが有用な情報を返した場合、その情報を使って詳細な回答を提供してください。**
**もしツールがエラーを返したり、関連情報が見つからなかった場合は、一般的な知識に基づいて回答し、情報が制限されていることを伝えてください。**
検索結果の情報源は常に明示してください。""",
tools=[bedrock_kb_tool],
)
これで、先ほど作成したBedrock KB検索ツールを使用するエージェントの基本構造が完成しました。
instructionの部分で、エージェントの動作方針を定義しています。
検索結果の構造と活用方法
Bedrock KBツールは以下のような構造の辞書を返します。
{
"status": str, # "success" または "error"
"results": [
{
"content": str, # 取得したコンテンツの本文
"source": str, # 情報のソース(S3のURI)
"score": float, # 検索結果のスコア(関連性の度合い)
"metadata": dict # 追加のメタデータ(属性と値のペア)
},
# ... 最大7件の結果が含まれる
],
"error_message": str # エラーが発生した場合のみ存在
}
実際に動かしてみた結果
agent1の上の階層に移動して、
実装したエージェントを起動するには、以下のコマンドを実行します。
uv run adk web
これでADKのウェブインターフェースが起動し、エージェントとチャットができるようになります。
実際に質問してみると、以下のような画面が表示されます。
回答の詳細情報を確認すると、以下のように表示されます。
上の画像を見ると、ちゃんとナレッジベースから情報を取得して回答に活用していることが分かります。特に「ツールを使用」の部分では、どのような情報をBedrock KBから取得したかが確認できます。
mastraと同じような感覚で、Google ADKはLangGraphよりもデバッグが楽ですよね。特にADKはエージェントの挙動が直感的で、状態のトラッキングやログが見やすく設計されているので、開発が楽になりそうだ。
まとめ
この記事では、Google ADKを使ってAmazon Bedrock Knowledge Baseをツールとして活用する方法を紹介しました。
Google製ということもあって(Vertex AI)などのGCPサービスとの連携は最初から用意されていて、使い勝手良さそうだと思いました。言語モデルの設定も。
retrieveの回答結果のフォーマットについてかなり迷いました。これが最適かどうかはまだ判断できませんが、とりあえず一旦形にしています。
より高度な検索結果を得たい場合は、Bedrockのフィルター設定やリランキング設定など、様々なパラメータを追加する必要がありそうです。
サンプルコード
作成したコードは以下のGitHubリポジトリで公開しています。興味のある方はぜひご参照ください。
登録データ
Amazon Bedrockナレッジベースには以下の資料を登録しています。
Discussion