TalentXの籔下です。
TalentXでは生成AIを用いた機能開発や業務効率化を進めており、その取り組みの中でAWS BedrockのナレッジベースのAPIであるRetrieve APIとRetrieveAndGenerate APIを、aws-sdk-go-v2経由で試してみました。
本記事では、S3にアップロードしたPDFの読み込み方法、事前に用意したベクトルデータベースを利用したRAG（Retrieval Augmented Generation）の使い方などを紹介します。

API概要
- Retrieve API
  - 実行してみる
  - Retrieve APIの実装
- RetrieveAndGenerate API
まとめ

API概要

ここでは、Retrieve API と RetrieveAndGenerate API について説明します。
また今回利用するaws-sdk-go-v2のbedrockagentruntimeに関する詳細は以下を参照ください。

pkg.go.dev

Retrieve API

Retrieve API は、ナレッジベースを活用して RAGの前段となるベクトル検索結果を返す API です。

docs.aws.amazon.com

まずは、PDF を S3 に格納し、そのバケットをデータソースとするナレッジベースを作成してから、その PDF に対して検索を実行してみます。サンプルとして、警視庁が公開している令和５年の犯罪情勢を利用しました。（特に深い意味はありません。笑）

実行してみる

まずはサンプル実装を動かし、結果を確認します。

package main

import (
    "context"
    "encoding/json"
    "fmt"
    "log"

    "github.com/aws/aws-sdk-go-v2/config"
    "github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime"
    "github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime/types"
    "github.com/aws/aws-sdk-go/aws"
)

func main() {

    cfg, err := config.LoadDefaultConfig(context.TODO())
    if err != nil {
        log.Fatalf("unable to load SDK config, %v", err)
    }

    client := bedrockagentruntime.New(bedrockagentruntime.Options{
        Region:      "ap-northeast-1",
        Credentials: cfg.Credentials,
    })

    input := &bedrockagentruntime.RetrieveInput{
        KnowledgeBaseId: aws.String("ナレッジベースのID"),
        RetrievalQuery: &types.KnowledgeBaseQuery{
            Text: aws.String("特殊詐欺による被害総額を教えてください"),
        },
    }

    output, err := client.Retrieve(context.TODO(), input)
    if err != nil {
        log.Fatalf("failed to retrieve and generate: %v", err)
    }
    outputJson, err := json.MarshalIndent(output, "", "  ")
    if err != nil {
        log.Fatalf("failed to marshal output: %v", err)
    }

    fmt.Println(string(outputJson))
}

実行結果は以下の通りです。

{
  "RetrievalResults": [
    {
      "Content": {
        "Text": "11     ウ 特殊詐欺      特殊詐欺については、事件の背後にいる暴力団や準暴力団を含む匿名・流動型犯     罪グループが、資金の供給、実行犯の周旋、犯行ツールの提供等を行い、犯行の分     業化と匿名化を図った上で、組織的に敢行している実態にあり、令和５年の認知件     数は１万 9,038 件、被害総額は約 452.6 億円と昨年に続き増加となり、深刻な情勢     が続いている（それぞれ前年比で 8.4％、22.0％増加）（図 15）。"
      },
      "Location": {
        "Type": "S3",
        "ConfluenceLocation": null,
        "S3Location": {
          "Uri": "s3://xxxxxxx(backet name)/r5_report_r.pdf"
        },
        "SalesforceLocation": null,
        "SharePointLocation": null,
        "WebLocation": null
      },
      "Metadata": {
        "x-amz-bedrock-kb-chunk-id": {},
        "x-amz-bedrock-kb-data-source-id": {},
        "x-amz-bedrock-kb-document-page-number": {},
        "x-amz-bedrock-kb-source-uri": {}
      },
      "Score": 0.6874807
    },

・・・

Retrieve API は、名前の通り検索結果を返す API なので、プロンプトに応じて根拠となるデータを返します。結果を見ると、「特殊詐欺については、〜〜被害総額は約 452.6 億円」といった記載があり、PDF 内にも同様の記述があることを確認できました。また、結果には Score が含まれていますが、これは 0〜1 の範囲で関連度を表し、数値が高いほどクエリとの関連性が高いことを示すようです。

Retrieve APIの実装

Retrieve APIは以下のようになっています。

func (c *Client) Retrieve(ctx context.Context, params *RetrieveInput, optFns ...func(*Options)) (*RetrieveOutput, error)

第二引数に RetrieveInput 構造体を渡して必要な情報を指定します。

type RetrieveInput struct {

    KnowledgeBaseId *string

    RetrievalQuery *types.KnowledgeBaseQuery

    NextToken *string

    RetrievalConfiguration *types.KnowledgeBaseRetrievalConfiguration

}

簡単に各フィールドについて説明します。

フィールド名	内容
KnowledgeBaseId	ナレッジベースID
RetrievalQuery	検索クエリを含む構造体
NextToken	ページネーション用のトークン
RetrievalConfiguration	クエリ実行方法や取得結果設定を含む構造体

今回はKnowledgeBaseIdとRetrievalQueryのみを使用しています。該当の実装箇所です。

input := &bedrockagentruntime.RetrieveInput{
    KnowledgeBaseId: aws.String("ナレッジベースのID"),
    RetrievalQuery: &types.KnowledgeBaseQuery{
        Text: aws.String("特殊詐欺による被害総額を教えてください"),
    },
}

KnowledgeBaseIdに指定するナレッジベースのIDはコンソール上で確認でき、RetrievalQueryに指定するtypes.KnowledgeBaseQueryは以下のような構造です。

type KnowledgeBaseQuery struct {
    Text *string
}

このTextにクエリをセットして渡すだけです。
なお、RetrievalConfigurationは今回は利用していませんが、KnowledgeBaseVectorSearchConfiguration構造体を使うことになります。これは後述の RetrieveAndGenerate API の説明時に詳細を記します。

RetrieveAndGenerate API

Retrieve API が検索結果のみを返すのに対し、RetrieveAndGenerate API は検索結果をもとに回答を生成する、いわゆる生成 AI 的な利用が可能な API です。また、RetrieveAndGenerate API では、S3 に格納したファイルを直接指定して回答を生成したり、ナレッジベースを利用して回答を生成したりすることができます。

docs.aws.amazon.com

S3上のファイルを直接指定

Retrieve API ではナレッジベース ID が必須だったため、PDF に対して検索する際には事前にナレッジベースを作成し、S3 バケットをデータソースとして同期する必要がありました。一方、RetrieveAndGenerate API では、S3 に格納したファイルを直接指定して会話（問い合わせ）が可能です。 Retrieve API を試したときに作成したナレッジベースやベクトルデータベースを削除した状態で、以下のコードを実行します。

package main

import (
    "context"
    "fmt"
    "log"

    "github.com/aws/aws-sdk-go-v2/aws"
    "github.com/aws/aws-sdk-go-v2/config"
    "github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime"
    "github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime/types"
)

func main() {

    cfg, err := config.LoadDefaultConfig(context.TODO())
    if err != nil {
        log.Fatalf("unable to load SDK config, %v", err)
    }

    client := bedrockagentruntime.New(bedrockagentruntime.Options{
        Region:      "ap-northeast-1",
        Credentials: cfg.Credentials,
    })

    documentS3URI := "s3://{バケット名}/r5_report_r.pdf"
    modelArn := "arn:aws:bedrock:ap-northeast-1::foundation-model/anthropic.claude-3-haiku-20240307-v1:0"
    question := "特殊詐欺による被害総額を教えてください"

    // パラメータ作成（
    input := &bedrockagentruntime.RetrieveAndGenerateInput{
        Input: &types.RetrieveAndGenerateInput{
            Text: &question,
        },
        RetrieveAndGenerateConfiguration: &types.RetrieveAndGenerateConfiguration{
            ExternalSourcesConfiguration: &types.ExternalSourcesRetrieveAndGenerateConfiguration{
                ModelArn: &modelArn,
                Sources: []types.ExternalSource{
                    {
                        SourceType: types.ExternalSourceTypeS3,
                        S3Location: &types.S3ObjectDoc{
                            Uri: &documentS3URI,
                        },
                    },
                },
            },
            Type: types.RetrieveAndGenerateTypeExternalSources,
        },
    }

    output, err := client.RetrieveAndGenerate(context.TODO(), input)
    if err != nil {
        log.Fatalf("failed to retrieve and generate: %v", err)
    }

    fmt.Println(aws.ToString(output.Output.Text))
}

実行結果は以下です。

$ go run ./main.go
特殊詐欺による被害総額は、令和5年が約452.6億円と、前年比で22.0%増加しており、深刻な情勢が続いています。

今回は RetrieveAndGenerate で返却される結果（生成された回答）のみを表示しましたが、ナレッジベースなしでも期待どおりの回答が得られました。

RetrieveAndGenerate APIの実装

RetrieveAndGenerate API は以下のようになってます。

func (c *Client) RetrieveAndGenerate(ctx context.Context, params *RetrieveAndGenerateInput, optFns ...func(*Options)) (*RetrieveAndGenerateOutput, error)

RetrieveAndGenerateInput

Retrieve API 同様、第二引数に必要な情報を渡す構造体が存在します。

type RetrieveAndGenerateInput struct {

    Input *types.RetrieveAndGenerateInput

    RetrieveAndGenerateConfiguration *types.RetrieveAndGenerateConfiguration

    SessionConfiguration *types.RetrieveAndGenerateSessionConfiguration

    SessionId *string
}

簡単に各フィールドについて説明します。

フィールド名	内容
Input	クエリ（質問や入力データ）を格納する構造体
RetrieveAndGenerateConfiguration	クエリや取得プロセスに関する設定を含む構造体
SessionConfiguration	セッションに関する設定を含む構造体
SessionId	セッションを識別する一意の ID

セッション周りに関しては継続的な会話に用いられますが、今回は割愛します。該当の実装箇所です。

    input := &bedrockagentruntime.RetrieveAndGenerateInput{
        Input: &types.RetrieveAndGenerateInput{
            Text: &question,
        },
        RetrieveAndGenerateConfiguration: &types.RetrieveAndGenerateConfiguration{
            ExternalSourcesConfiguration: &types.ExternalSourcesRetrieveAndGenerateConfiguration{
                ModelArn: &modelArn,
                Sources: []types.ExternalSource{
                    {
                        SourceType: types.ExternalSourceTypeS3,
                        S3Location: &types.S3ObjectDoc{
                            Uri: &documentS3URI,
                        },
                    },
                },
            },
            Type: types.RetrieveAndGenerateTypeExternalSources,
        },
    }

Input *types.RetrieveAndGenerateInput

RetrieveAndGenerateInput は bedrockagentruntime パッケージで定義されており、その中で参照する types.RetrieveAndGenerateInput は以下のような構造です。

type RetrieveAndGenerateInput struct {
    Text *string
}

Text にクエリをセットする点は、Retrieve API の KnowledgeBaseQuery と類似しています。

RetrieveAndGenerateConfiguration *types.RetrieveAndGenerateConfiguration

RetrieveAndGenerateConfiguration は以下のような構造です。

type RetrieveAndGenerateConfiguration struct {

    Type RetrieveAndGenerateType

    ExternalSourcesConfiguration *ExternalSourcesRetrieveAndGenerateConfiguration

    KnowledgeBaseConfiguration *KnowledgeBaseRetrieveAndGenerateConfiguration
}

Type フィールドには、ナレッジベースを使う場合は "KNOWLEDGE_BASE"、外部リソースを使う場合は "EXTERNAL_SOURCES" を指定できます。またこれらは以下のようにEnumとして定義されています。

const (
    RetrieveAndGenerateTypeKnowledgeBase   RetrieveAndGenerateType = "KNOWLEDGE_BASE"
    RetrieveAndGenerateTypeExternalSources RetrieveAndGenerateType = "EXTERNAL_SOURCES"
)

今回は S3 を利用するため RetrieveAndGenerateTypeExternalSources を指定しました。

ExternalSourcesConfiguration または KnowledgeBaseConfiguration のいずれかを、Type に応じて設定します。今回は S3 利用のため、ExternalSourcesConfiguration を指定します。 ExternalSourcesConfiguration に定義される ExternalSourcesRetrieveAndGenerateConfiguration は以下の通りです。

type ExternalSourcesRetrieveAndGenerateConfiguration struct {

    ModelArn *string

    Sources []ExternalSource

    GenerationConfiguration *ExternalSourcesGenerationConfiguration
}

ModelArn には利用する言語モデル（LLM）を指定します。利用可能なモデルは下記を参照してください。

docs.aws.amazon.com

Sources には ExternalSource 構造体をリストで指定します。 ExternalSource は以下のようになっています。

type ExternalSource struct {

    SourceType ExternalSourceType

    ByteContent *ByteContentDoc

    S3Location *S3ObjectDoc
}

SourceTypeには"S3"か"BYTE_CONTENT"を指定可能です。またこちらもEnumで定義されています。

type ExternalSourceType string

// Enum values for ExternalSourceType
const (
    ExternalSourceTypeS3          ExternalSourceType = "S3"
    ExternalSourceTypeByteContent ExternalSourceType = "BYTE_CONTENT"
)

ByteContent と S3Location は SourceType に指定する値によってどちらかを指定します。今回はS3Location を指定します。 S3Location に定義されている S3ObjectDoc は以下の通りです。

type S3ObjectDoc struct {

    Uri *string
}

S3 に格納したファイルの URI を指定することで、そのファイルを元に回答を生成します。

GenerationConfiguration ではモデル推論用のカスタムパラメータ、ガードレール、プロンプトテンプレートなどを設定可能ですが、今回は説明を割愛します。

ナレッジベースを指定

同じ PDF の例ばかりでは飽きるので、次はメタデータフィルタリングと組み合わせて、ナレッジベースを指定して RetrieveAndGenerate API を試します。まず、以下のような形式の CSV ファイルを S3 にアップロードします。

さらに、csvファイル名.metadata.json というファイル名で同じバケットに JSON ファイルをアップロードします。

{
    "metadataAttributes": {
        "table_name": "items"
    },
    "documentStructureConfiguration": {
        "type": "RECORD_BASED_STRUCTURE_METADATA",
        "recordBasedStructureMetadata": {
            "contentFields": [
                { "fieldName": "name" }
            ],
            "metadataFieldsSpecification": {
                "fieldsToInclude": [
                    { "fieldName": "item_id" },
                    { "fieldName": "category_id" }
                ]
            }
        }
    }
}

このメタデータフィルタリングを用いて category_id が 2 の商品情報のみを取得するコードは以下です。

package main

import (
    "context"
    "fmt"
    "log"

    "github.com/aws/aws-sdk-go-v2/aws"
    "github.com/aws/aws-sdk-go-v2/config"
    "github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime"
    "github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime/document"
    "github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime/types"
)

func main() {

    cfg, err := config.LoadDefaultConfig(context.TODO())
    if err != nil {
        log.Fatalf("unable to load SDK config, %v", err)
    }

    client := bedrockagentruntime.New(bedrockagentruntime.Options{
        Region:      "ap-northeast-1",
        Credentials: cfg.Credentials,
    })

    modelArn := "arn:aws:bedrock:ap-northeast-1::foundation-model/anthropic.claude-3-haiku-20240307-v1:0"
    question := "商品を教えてください"

    filter := &types.RetrievalFilterMemberEquals{
        Value: types.FilterAttribute{
            Key:   aws.String("category_id"),
            Value: document.NewLazyDocument("2"),
        },
    }

    input := &bedrockagentruntime.RetrieveAndGenerateInput{
        Input: &types.RetrieveAndGenerateInput{
            Text: &question,
        },
        RetrieveAndGenerateConfiguration: &types.RetrieveAndGenerateConfiguration{
            Type: types.RetrieveAndGenerateTypeKnowledgeBase,
            KnowledgeBaseConfiguration: &types.KnowledgeBaseRetrieveAndGenerateConfiguration{
                KnowledgeBaseId: aws.String("ナレッジベースのID"),

                RetrievalConfiguration: &types.KnowledgeBaseRetrievalConfiguration{
                    VectorSearchConfiguration: &types.KnowledgeBaseVectorSearchConfiguration{
                        Filter: filter,
                    },
                },

                ModelArn: aws.String(modelArn),
            },
        },
    }

    output, err := client.RetrieveAndGenerate(context.TODO(), input)
    if err != nil {
        log.Fatalf("failed to retrieve and generate: %v", err)
    }

    fmt.Println(aws.ToString(output.Output.Text))
}

実行結果：

$ go run ./main.go
この検索結果には、スマートフォンとスマートウォッチという2つの商品が含まれています。

category_id="2"を指定した結果、該当する商品のみが回答に含まれました。今回は RetrieveAndGenerateInput にナレッジベース関連の設定やメタデータフィルタリングを指定しています。該当の実装です。

   input := &bedrockagentruntime.RetrieveAndGenerateInput{
        Input: &types.RetrieveAndGenerateInput{
            Text: &question,
        },
        RetrieveAndGenerateConfiguration: &types.RetrieveAndGenerateConfiguration{
            Type: types.RetrieveAndGenerateTypeKnowledgeBase,
            KnowledgeBaseConfiguration: &types.KnowledgeBaseRetrieveAndGenerateConfiguration{
                KnowledgeBaseId: aws.String("ナレッジベースのID"),

                RetrievalConfiguration: &types.KnowledgeBaseRetrievalConfiguration{
                    VectorSearchConfiguration: &types.KnowledgeBaseVectorSearchConfiguration{
                        Filter: filter,
                    },
                },

                ModelArn: aws.String(modelArn),
            },
        },
    }

RetrieveAndGenerateConfiguration *types.RetrieveAndGenerateConfiguration

ナレッジベースを利用する場合は RetrieveAndGenerateConfiguration で Type に RetrieveAndGenerateTypeKnowledgeBase を指定し、KnowledgeBaseConfiguration を設定します。再掲ですがRetrieveAndGenerateConfigurationの実装です。

type RetrieveAndGenerateConfiguration struct {

    Type RetrieveAndGenerateType

    ExternalSourcesConfiguration *ExternalSourcesRetrieveAndGenerateConfiguration

    KnowledgeBaseConfiguration *KnowledgeBaseRetrieveAndGenerateConfiguration

}

KnowledgeBaseRetrieveAndGenerateConfiguration

KnowledgeBaseRetrieveAndGenerateConfiguration は以下のようになっています。

type KnowledgeBaseRetrieveAndGenerateConfiguration struct {

    KnowledgeBaseId *string

    ModelArn *string

    GenerationConfiguration *GenerationConfiguration

    OrchestrationConfiguration *OrchestrationConfiguration

    RetrievalConfiguration *KnowledgeBaseRetrievalConfiguration

}

KnowledgeBaseId や ModelArn は Retrieve API と同様に指定します。GenerationConfiguration や OrchestrationConfiguration は必要に応じて、回答生成時のパラメータやクエリ変換・会話履歴考慮などを設定します。今回はConfigurationに関してRetrievalConfiguration のみを利用しました。RetrievalConfiguration は KnowledgeBaseRetrievalConfiguration 構造体であり、KnowledgeBaseVectorSearchConfiguration を内部で利用します。

type KnowledgeBaseRetrievalConfiguration struct {

    VectorSearchConfiguration *KnowledgeBaseVectorSearchConfiguration

}

type KnowledgeBaseVectorSearchConfiguration struct {

    Filter RetrievalFilter

    NumberOfResults *int32

    OverrideSearchType SearchType

}

Filterの該当実装を改めて見てください。

filter := &types.RetrievalFilterMemberEquals{
    Value: types.FilterAttribute{
        Key:   aws.String("category_id"),
        Value: document.NewLazyDocument("2"),
    },
}

~~
RetrievalConfiguration: &types.KnowledgeBaseRetrievalConfiguration{
    VectorSearchConfiguration: &types.KnowledgeBaseVectorSearchConfiguration{
        Filter: filter,
    },
},
~~

このFilter を利用することでメタデータフィルタが可能です。 aws-sdk-go(v1)では RetrievalFilter は構造体として定義されていましたが、v2 では以下のようにインターフェースとして定義されています。

type RetrievalFilter interface {
    isRetrievalFilter()
}

そのため、このインターフェースを満たす構造体を指定する必要があります。今回の例ではcategory_id="2"を指定してフィルタしており、このイコールを表現するRetrievalFilterMemberEqualsを利用しています。以下の表に、利用可能なフィルタをまとめました。

構造体	内容
RetrievalFilterMemberEquals	Value FilterAttribute\| Key = Valueでフィルタ
RetrievalFilterMemberGreaterThan	Key > Valueでフィルタ
RetrievalFilterMemberGreaterThanOrEquals	Key >= Valueでフィルタ
RetrievalFilterMemberIn	Keyの値がValue（配列）のいずれかに一致
RetrievalFilterMemberLessThan	Key < Valueでフィルタ
RetrievalFilterMemberLessThanOrEquals	Key <= Valueでフィルタ
RetrievalFilterMemberListContains	Keyがリスト型属性で、そのリストがValueを含んでいる場合にフィルタ
RetrievalFilterMemberNotEquals	Key ≠ Valueでフィルタ
RetrievalFilterMemberNotIn	Keyの値がValue（配列）のいずれにも一致しない場合にフィルタ
RetrievalFilterMemberStartsWith	Key の値が Value で始まるものをフィルタ（Amazon OpenSearch Serverless vector store 限定）
RetrievalFilterMemberStringContains	Key の値が文字列またはリスト型属性で、その中の文字列に Value が含まれる場合にフィルタ

またこれらのフィルタを複数設定したい場合はValue []RetrievalFilter をフィールドに持つ以下の構造体を利用することで実現可能です。

構造体	内容
RetrievalFilterMemberAndAll	複数のフィルタをAnd条件でフィルタする
RetrievalFilterMemberOrAll	複数のフィルタをOr条件でフィルタする

以下のようなイメージです。

itemIDFilter := &types.RetrievalFilterMemberEquals{
    Value: types.FilterAttribute{
        Key:   aws.String("item_id"),
        Value: document.NewLazyDocument("1"),
    },
}

categoryIDFilter := &types.RetrievalFilterMemberEquals{
    Value: types.FilterAttribute{
        Key:   aws.String("category_id"),
        Value: document.NewLazyDocument("2"),
    },
}

filter := &types.RetrievalFilterMemberAndAll{
    Value: []types.RetrievalFilter{
        itemIDFilter,
        categoryIDFilter,
    },
}

まとめ

本記事では、Bedrock のナレッジベースを利用した Retrieve API および RetrieveAndGenerate API を、aws-sdk-go-v2 で実際に触ってみました。S3 上の PDF の読み込み、ナレッジベースやベクトルデータベースを用いた RAG、メタデータフィルタリングによる絞り込みなどの手順を紹介しました。

またエンジニア採用も進めていますので、是非ご応募お待ちしています！
私とのカジュアル面談も受け付けていますので気軽にご応募ください！

i-myrefer.jp

TalentX Tech Blog

Tech Blog

aws-sdk-go-v2で試すAWS Bedrockナレッジベース：Retrieve & RetrieveAndGenerate API活用

API概要

Retrieve API

実行してみる

Retrieve APIの実装

RetrieveAndGenerate API

S3上のファイルを直接指定

RetrieveAndGenerate APIの実装

RetrieveAndGenerateInput

Input *types.RetrieveAndGenerateInput

RetrieveAndGenerateConfiguration *types.RetrieveAndGenerateConfiguration

ナレッジベースを指定

RetrieveAndGenerateConfiguration *types.RetrieveAndGenerateConfiguration

KnowledgeBaseRetrieveAndGenerateConfiguration

まとめ