Amazon Bedrock Knowledge Base MCP Server

Instructions

Implementation Reference

• src/bedrock_kb_mcp_server/main.py:858-897 (handler)

  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

• src/bedrock_kb_mcp_server/bedrock_client.py:541-598 (handler)

  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,  # 実行したクエリテキスト（確認用）
          }