👻

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のウェブインターフェースが起動し、エージェントとチャットができるようになります。

実際に質問してみると、以下のような画面が表示されます。

エージェント回答画面

回答の詳細情報を確認すると、以下のように表示されます。

回答詳細1

回答詳細2

上の画像を見ると、ちゃんとナレッジベースから情報を取得して回答に活用していることが分かります。特に「ツールを使用」の部分では、どのような情報をBedrock KBから取得したかが確認できます。

mastraと同じような感覚で、Google ADKはLangGraphよりもデバッグが楽ですよね。特にADKはエージェントの挙動が直感的で、状態のトラッキングやログが見やすく設計されているので、開発が楽になりそうだ。

まとめ

この記事では、Google ADKを使ってAmazon Bedrock Knowledge Baseをツールとして活用する方法を紹介しました。

Google製ということもあって(Vertex AI)などのGCPサービスとの連携は最初から用意されていて、使い勝手良さそうだと思いました。言語モデルの設定も。

retrieveの回答結果のフォーマットについてかなり迷いました。これが最適かどうかはまだ判断できませんが、とりあえず一旦形にしています。

より高度な検索結果を得たい場合は、Bedrockのフィルター設定やリランキング設定など、様々なパラメータを追加する必要がありそうです。

サンプルコード

作成したコードは以下のGitHubリポジトリで公開しています。興味のある方はぜひご参照ください。
https://github.com/Kota8102/google-adk-bedrock-kb-sample

登録データ

Amazon Bedrockナレッジベースには以下の資料を登録しています。

参考リソース

Discussion