retrieve
Execute RAG queries against Amazon Bedrock Knowledge Bases to retrieve relevant documents using vector search for enhanced information retrieval.
Instructions
Knowledge Baseに対してRAG（Retrieval-Augmented Generation）クエリを実行します。
ベクトル検索を使用して、クエリに関連するドキュメントを取得します。
Args: knowledge_base_id: クエリ対象のKnowledge BaseのID query: 検索クエリのテキスト number_of_results: 返す結果の数（デフォルト: 5、範囲: 1-100）
Returns: RetrieveResponseDict: クエリ結果 - results: 検索結果のリスト（各結果にはcontent、location、score、metadataが含まれる） - query: 実行したクエリテキスト
Raises: ValueError: 入力値が無効な場合（knowledge_base_idやqueryが空、number_of_resultsが範囲外など）
Input Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| knowledge_base_id | Yes | ||
| query | Yes | ||
| number_of_results | No |
Output Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| query | Yes | ||
| results | Yes |
Implementation Reference
- The 'retrieve' MCP tool handler function registered in the MCP server. It validates inputs and delegates the actual RAG query to the bedrock_client.
def retrieve( knowledge_base_id: str, query: str, number_of_results: int = 5 ) -> RetrieveResponseDict: """ Knowledge Baseに対してRAG（Retrieval-Augmented Generation）クエリを実行します。 ベクトル検索を使用して、クエリに関連するドキュメントを取得します。 Args: knowledge_base_id: クエリ対象のKnowledge BaseのID query: 検索クエリのテキスト number_of_results: 返す結果の数（デフォルト: 5、範囲: 1-100） Returns: RetrieveResponseDict: クエリ結果 - results: 検索結果のリスト（各結果にはcontent、location、score、metadataが含まれる） - query: 実行したクエリテキスト Raises: ValueError: 入力値が無効な場合（knowledge_base_idやqueryが空、number_of_resultsが範囲外など） """ # 入力値のバリデーション（共通関数を使用） # knowledge_base_idとqueryは必須パラメータです knowledge_base_id = validate_required_string(knowledge_base_id, "knowledge_base_id") query = validate_required_string(query, "query") # number_of_resultsは1から100の範囲で指定する必要があります # AWS APIの制限に従います if number_of_results < 1 or number_of_results > 100: raise ValueError("number_of_results must be between 1 and 100") # Bedrockクライアントを使用してRAGクエリを実行 # ベクトル検索を使用して、クエリに関連するドキュメントを取得します # 結果は関連度スコアでソートされ、指定された数の結果が返されます result = bedrock_client.retrieve( knowledge_base_id, # 前後の空白は既に削除済み query, # 前後の空白は既に削除済み number_of_results # 返す結果の数 ) return result - The bedrock_client's implementation of 'retrieve', which interacts with the Bedrock Agent Runtime API to perform the actual vector search.
def retrieve( self, knowledge_base_id: str, query: str, number_of_results: int = 5, ) -> RetrieveResponseDict: """ Knowledge Baseに対してRAG（Retrieval-Augmented Generation）クエリを実行します。 ベクトル検索を使用して、クエリに関連するドキュメントを取得します。 Args: knowledge_base_id: クエリ対象のKnowledge BaseのID query: 検索クエリのテキスト number_of_results: 返す結果の数（デフォルト: 5、最大: 100） Returns: RetrieveResponseDict: クエリ結果 - results: 検索結果のリスト 各結果には以下の情報が含まれます: - content: ドキュメントの内容 - location: ドキュメントの場所（S3 URIなど） - score: 関連度スコア - metadata: メタデータ - query: 実行したクエリテキスト Raises: ClientError: AWS API呼び出しが失敗した場合 """ try: # Bedrock Agent Runtime APIを使用してRAGクエリを実行 # retrieve APIは、Knowledge Baseに対してベクトル検索を実行し、 # クエリに関連するドキュメントチャンクを返します response = self.bedrock_agent_runtime.retrieve( knowledgeBaseId=knowledge_base_id, # クエリ対象のKnowledge BaseのID retrievalConfiguration={ "vectorSearchConfiguration": { "numberOfResults": number_of_results, # 返す結果の数（1-100の範囲） # オプション: "overrideSearchType": "HYBRID" # ハイブリッド検索（将来の拡張） # オプション: "filter": {...} # メタデータフィルター（将来の拡張） } }, retrievalQuery={"text": query}, # 検索クエリのテキスト # オプション: nextToken: ページネーション用のトークン（大量の結果がある場合） ) # 取得した結果の数をログに記録 # 検索結果の数をログに記録します（デバッグや監視に有用） retrieval_results = response.get("retrievalResults", []) logger.info(f"Retrieved {len(retrieval_results)} results") # レスポンスを整形して返す # AWS APIのレスポンスから検索結果を抽出し、クエリテキストと一緒に返します # 各結果には、content（テキスト内容）、location（S3 URIなど）、score（関連度）、metadataが含まれます return { "results": retrieval_results, # 検索結果のリスト（関連度順にソート済み） "query": query, # 実行したクエリテキスト（確認用） }