🎃

AWS CDKでPineconeとBedrockのKnowledgeBaseとAgent環境を構築してみた

2024/06/16に公開

Agent for Bedrockを使って、AIエージェントをいくつか作ってみたいと思い、まずは環境をパッと構築できるようにAWS CDKでIaCを実装しました。AWS CDKを使ってPineconeとBedrockのKnowledgeBaseとAgentの環境を構築してみたので、その内容をご紹介します。詳細な設定やコードは、以下のリポジトリにて確認できます。なお、今回紹介していないOpensearch Serverlessを使用するバージョンも含まれています。

https://github.com/memememomo/bedrock-iac

以下は、今回構築したシステムの構成図です。

Pineconeとは

https://www.pinecone.io/

Pineconeは、高性能なベクトル検索アプリケーションを簡単に構築できるフルマネージド型のクラウドネイティブなベクトルデータベースサービスです。主な特徴は以下のとおりです。

  • 数十億規模のベクトルデータでも高速検索が可能
  • データの追加・更新に合わせてインデックスがリアルタイムに更新される
  • ベクトル検索とメタデータフィルタの組み合わせで適切な結果を素早く取得可能
  • 運用の手間いらずで、導入・利用・拡張が容易

AWS Marketplace経由でPineconeと連携

AWSでは、Marketplaceを通じてPinecone serverlessをサブスクライブすることができます。このようにすることで、AWSアカウントと連携してPineconeを利用でき、利用料金もAWSの請求と一元化されるようになります。

AWS Marketplaceで、Pineconeをサブスクライブすると、以下のような画面になります。

MarketplaceでPineconeをサブスクライブ

Pay as you goという表記があるとおり、Pineconeの料金は使った分だけになります。

この画面からPineconeのサイトに行くと、アカウント作成画面が表示されるので、アカウントがない場合は作成しておきます。アカウントを作成すると、Pineconeの設定画面に入ることができます。

後ほどAPIキーを使用するので、以下のページに移動して、控えておきます。

PineconeのAPIキー

Pinecone情報を保存・参照するためのシークレットを作成する

このあとの処理で、PineconeのAPIキーを参照する必要があるため、SecretsManagerに保存しておきます。AWS CDKで、APIキー用の入れ物をSecretsManagerに作るコードを記述します。

import { CfnOutput, RemovalPolicy, SecretValue, Stack, StackProps } from "aws-cdk-lib";
import { Secret } from "aws-cdk-lib/aws-secretsmanager";
import { Construct } from "constructs";
import { Config } from "../service/types";
import { EXPORT_NAME, PARAMS } from "../service/const";

// PineconeSecretStackPropsインターフェイスを定義
export interface PineconeSecretStackProps extends StackProps {
  config: Config; // 設定情報を含むプロパティ
}

// PineconeSecretStackクラスを定義
export class PineconeSecretStack extends Stack {
  apiKeySecret: Secret; // APIキーのシークレット

  // コンストラクタ
  constructor(scope: Construct, id: string, props: PineconeSecretStackProps) {
    super(scope, id, props); // 親クラスStackのコンストラクタを呼び出し

    const { prefix } = props.config; // コンフィグからプレフィックスを取得

    // Secrets Managerで新しいシークレットを作成
    this.apiKeySecret = new Secret(this, "PineconeApiKeySecret", {
      secretName: PARAMS.SECRET_KEY.PINECONE_API_KEY(prefix), // シークレット名を設定
      removalPolicy: RemovalPolicy.DESTROY, // 削除ポリシーを設定(スタック削除時にシークレットも削除)
      secretObjectValue: {
        apiKey: new SecretValue("Pinecone API Key"), // シークレットの値を設定
      },
    });

    // シークレットのARNをCloudFormation出力として定義
    new CfnOutput(this, "PineconeApiKeySecretArn", {
      exportName: EXPORT_NAME.PINECONE_SECRET, // エクスポート名を設定
      value: this.apiKeySecret.secretArn, // シークレットのARNを出力
    });
  }
}

実行コマンド例は以下のとおりです(必要に応じて環境変数などを設定して、デプロイ先を指定する)。

$ cdk deploy PineconeSecretStack

これにより、SecretsManagerにAPIキーを入力する入れ物が作成されるので、AWSコンソールからAPIキーを設定します。

今後の処理でAPIキーが必要になった場合は、ここで作成されたシークレットを参照します。

Pineconeのインデックスを作成するカスタムリソースを実装する

次にPineconeのインデックスを作成します。
PineconeはAWSのリソースではないため、AWS CDKで直接インデックスを作成することはできません。
ですので、カスタムリソースを利用します。カスタムリソースのLambdaの中では、Pineconeのライブラリを使用して、インデックスを作成します。

今回は、Lambdaの実装としてGo言語を使用しました。
以下のコードは、 github.com/pinecone-io/go-pinecone を利用して、Pineconeのインデックスを作成しているものになります。

// Pineconeクライアントを新規作成
func (p *PineconeIndex) newClient() (*pinecone.Client, error) {
	// クライアントパラメータを設定して新しいクライアントを作成
	client, err := pinecone.NewClient(pinecone.NewClientParams{
		ApiKey: p.apiKey,
	})
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return client, nil
}

// インデックスを作成
func (p *PineconeIndex) CreateIndex(ctx context.Context) (string, map[string]interface{}, error) {
	// 新しいクライアントを作成
	client, err := p.newClient()
	if err != nil {
		return "", nil, errors.WithStack(err)
	}

	// クライアントを使用してインデックスを作成
	idx, err := client.CreateServerlessIndex(ctx, &pinecone.CreateServerlessIndexRequest{
		Name:      p.indexName,
		Dimension: p.dimension,
		Metric:    pinecone.Cosine,
		Cloud:     pinecone.Aws,
		Region:    p.region,
	})
	if err != nil {
		return "", nil, errors.WithStack(err)
	}

	// インデックスエンドポイントを保存
	err = p.saveIndexEndpoint(ctx, idx.Host)
	if err != nil {
		return "", nil, errors.WithStack(err)
	}

	return "PineconeIndex", nil, nil
}

作成したインデックスのエンドポイント情報は、SecretsManagerに保存するようにしています。

// インデックスエンドポイントをSecrets Managerに保存
func (p *PineconeIndex) saveIndexEndpoint(ctx context.Context, indexEndpoint string) error {
	// 新しいSecrets Managerクライアントを作成
	svc, err := newSecretsManagerClient(ctx, p.region)
	if err != nil {
		return errors.WithStack(err)
	}

	// シークレットを作成してエンドポイントを保存
	_, err = svc.CreateSecret(ctx, &secretsmanager.CreateSecretInput{
		Name:         aws.String(p.indexEndpointSecretKey),
		SecretString: aws.String(indexEndpoint),
	})
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

以上のLambdaのコードをカスタムリソースとして実行する設定を行っているのが、次のコードとなります。

import { Construct } from "constructs";
import * as lambda from "aws-cdk-lib/aws-lambda";
import * as ecr_assets from "aws-cdk-lib/aws-ecr-assets";
import * as iam from "aws-cdk-lib/aws-iam";
import { Provider } from "aws-cdk-lib/custom-resources";
import { RetentionDays } from "aws-cdk-lib/aws-logs";
import { CustomResource, Duration } from "aws-cdk-lib";
import { region, secretsManagerPartialArn } from "../service/util";

// PineconeProps型を定義
type PineconeProps = {
  apiKeySecretKey: string;          // APIキーのシークレットキー
  indexEndpointSecretKey: string;   // インデックスエンドポイントのシークレットキー
  dimension: number;                // インデックスの次元数
  indexName: string;                // インデックス名
};

// Pineconeクラスを定義
export class Pinecone extends Construct {
  constructor(scope: Construct, id: string, props: PineconeProps) {
    super(scope, id); // 親クラスConstructのコンストラクタを呼び出し

    // Lambda関数用のIAMロールを作成
    const customResourceRole = new iam.Role(this, "PineconeCustomResourceRole", {
      assumedBy: new iam.ServicePrincipal("lambda.amazonaws.com"), // Lambdaサービスプリンシパル
      inlinePolicies: {
        ssmPolicy: new iam.PolicyDocument({
          statements: [
            // Secrets Managerからシークレット値を取得するためのポリシーステートメント
            new iam.PolicyStatement({
              resources: [secretsManagerPartialArn(this, props.apiKeySecretKey)],
              actions: ["secretsmanager:GetSecretValue"],
            }),
            // Secrets Managerでシークレットを作成および削除するためのポリシーステートメント
            new iam.PolicyStatement({
              resources: [secretsManagerPartialArn(this, props.indexEndpointSecretKey)],
              actions: ["secretsmanager:CreateSecret", "secretsmanager:DeleteSecret"],
            }),
          ],
        }),
      },
      managedPolicies: [
        // Lambdaの基本的な実行権限を含むマネージドポリシーをアタッチ
        iam.ManagedPolicy.fromAwsManagedPolicyName("service-role/AWSLambdaBasicExecutionRole"),
      ],
    });

    // Dockerイメージを使用してLambda関数を作成
    const customResourceF = new lambda.DockerImageFunction(this, "CustomResourceFunction", {
      code: lambda.DockerImageCode.fromImageAsset("./src/resources/pinecone", {
        platform: ecr_assets.Platform.LINUX_ARM64, // ARM64プラットフォームを指定
      }),
      architecture: lambda.Architecture.ARM_64, // ARM64アーキテクチャを指定
      timeout: Duration.minutes(15), // タイムアウトを15分に設定
      role: customResourceRole, // 上で作成したIAMロールを使用
    });

    // カスタムリソースプロバイダーを作成
    const provider = new Provider(this, "PineconeProvider", {
      onEventHandler: customResourceF, // イベントハンドラとして上記のLambda関数を使用
      logRetention: RetentionDays.ONE_WEEK, // ログ保持期間を1週間に設定
    });

    // カスタムリソースを作成
    new CustomResource(this, "PineconeCustomResource", {
      serviceToken: provider.serviceToken, // プロバイダーのサービストークンを設定
      properties: {
        region: region(this), // リージョンを設定
        indexName: props.indexName, // インデックス名を設定
        dimension: props.dimension, // 次元数を設定
        apiKeySecretName: props.apiKeySecretKey, // APIキーシークレット名を設定
        indexEndpointSecretName: props.indexEndpointSecretKey, // インデックスエンドポイントシークレット名を設定
      },
    });
  }
}

実行コマンド例は以下のとおりです。

$ cdk deploy PineconeStack

実行が完了すると、Pineconeの管理画面でインデックスが作成されていることが確認できます。

Pineconeのインデックス

KnowledgeBaseを作成する

次にKnowledgeBaseを作成します。

KnowledgeBaseには、先程作成したPineconeの情報を設定する形になります。

// KnowledgeBase作成
// Pineconeのインデックスエンドポイントを取得
const endpoint = SecretValue.secretsManager(props.pineconeParams?.indexEndpointSecretKey!).unsafeUnwrap();

// BedrockKnowledgeBaseのCloudFormationリソースを作成
const cfnKnowledgeBase = new bedrock.CfnKnowledgeBase(this, "BedrockKnowledgeBase", {
  name: props.knowledgeBaseParams.name, // 知識ベースの名前を設定
  roleArn: this.knowledgeBaseRole.roleArn, // 知識ベースのロールARNを設定
  knowledgeBaseConfiguration: {
    type: "VECTOR", // 知識ベースのタイプをベクトルに設定
    vectorKnowledgeBaseConfiguration: {
      embeddingModelArn: embeddingModelArn(this, props.embeddingModelName), // 埋め込みモデルのARNを設定
    },
  },
  storageConfiguration: {
    type: "PINECONE", // ストレージのタイプをPineconeに設定
    pineconeConfiguration: {
      connectionString: `https://${endpoint}/`, // Pineconeの接続文字列を設定
      credentialsSecretArn: props.pineconeParams?.apiKeySecret.secretFullArn!, // PineconeのAPIキーシークレットのARNを設定
      fieldMapping: {
        metadataField: "metadata", // メタデータフィールドを設定
        textField: "text", // テキストフィールドを設定
      },
    },
  },
});

次に、DataSourceを作成し、Pineconeに入れるデータの基となるファイルを保存するためのS3バケットを設定します。

// BedrockDataSourceのCloudFormationリソースを作成
const cfnDataSource = new bedrock.CfnDataSource(this, "BedrockDataSource", {
  name: props.knowledgeBaseParams.name, // データソースの名前を設定
  dataSourceConfiguration: {
    s3Configuration: {
      bucketArn: props.knowledgeBaseParams.bucket.bucketArn, // S3バケットのARNを設定
    },
    type: "S3", // データソースのタイプをS3に設定
  },
  knowledgeBaseId: cfnKnowledgeBase.attrKnowledgeBaseId, // 知識ベースのIDを設定

  dataDeletionPolicy: "DELETE", // データ削除ポリシーを設定
  serverSideEncryptionConfiguration: {
    kmsKeyArn: key.keyArn, // サーバーサイド暗号化のKMSキーARNを設定
  },
  vectorIngestionConfiguration: {
    chunkingConfiguration: {
      chunkingStrategy: "FIXED_SIZE", // 固定サイズのチャンクング戦略を設定
      fixedSizeChunkingConfiguration: {
        maxTokens: 300, // 最大トークン数を300に設定
        overlapPercentage: 20, // オーバーラップ率を20%に設定
      },
    },
  },
});

実行コマンド例は以下のとおりです

$ cdk deploy BedrockStack

実行が完了すると、AWSコンソールでKnowledgeBaseが作成されていることが確認できます。

Agentを作成する

最後に、Agentを作成します。
Agentの設定は、ひとまずダミーとなる設定を入れておきます。
Agentを実装する際は、ダミーとなる設定を差し替えていくことを想定しています。

Agentでは、アクションとなる設定をOpenAPI形式で設定します。
OpenAPI形式で書かれたJsonファイルをS3バケットに保存しておき、Agentがそれを読み込んでアクションを実行する形になります。

まずは、アクションの設定を保存するためのS3バケットとダミーとなる設定ファイルをアップロードする処理を書きます。

import * as cdk from "aws-cdk-lib";
import { Construct } from "constructs";
import { BedrockAgentSchemaBucket } from "../construct/bedrock-agent-schema-bucket";
import { Config } from "../service/types";

// BedrockAgentSchemaStackPropsインターフェイスを定義
export interface BedrockAgentSchemaStackProps extends cdk.StackProps {
  config: Config; // 設定情報を含むプロパティ
}

// BedrockAgentSchemaStackクラスを定義
export class BedrockAgentSchemaStack extends cdk.Stack {
  public readonly schemaBucket: BedrockAgentSchemaBucket; // スキーマバケットの公開プロパティ

  // コンストラクタ
  constructor(scope: Construct, id: string, props: BedrockAgentSchemaStackProps) {
    super(scope, id); // 親クラスStackのコンストラクタを呼び出し

    // BedrockAgentSchemaBucketを作成し、schemaBucketプロパティに設定
    this.schemaBucket = new BedrockAgentSchemaBucket(this, "BedrockAgentSchemaBucket", {
      prefix: props.config.prefix, // コンフィグからプレフィックスを設定
      schemaFilePath: "./src/resources/schema", // スキーマファイルのパスを設定
    });
  }
}

以上の処理は、Agentを作成するスタックとは別スタックにしています。ダミー設定ファイルのアップロード処理の後に、Agentを作成する処理が実行されるようにするためです。

Agentを作成するためのコードは以下のようになります。

import * as cdk from "aws-cdk-lib";
import { Construct } from "constructs";
import { APP_NAME, EXPORT_NAME, PARAMS } from "../service/const";
import { BedrockAgent } from "../construct/bedrock-agent";
import { BedrockAgentSchemaBucket } from "../construct/bedrock-agent-schema-bucket";

// AgentStackPropsインターフェイスを定義
export interface AgentStackProps extends cdk.StackProps {
  schemaBucket: BedrockAgentSchemaBucket; // スキーマバケットを含むプロパティ
}

// AgentStackクラスを定義
export class AgentStack extends cdk.Stack {
  // コンストラクタ
  constructor(scope: Construct, id: string, props: AgentStackProps) {
    super(scope, id, props); // 親クラスStackのコンストラクタを呼び出し

    // 既存のCloudFormation出力からKnowledgeBase IDをインポート
    const knowledgeBaseId = cdk.Fn.importValue(EXPORT_NAME.KNOWLEDGE_BASE_ID);

    // BedrockAgentを作成
    const agent = new BedrockAgent(this, "BedrockAgent", {
      prefix: APP_NAME, // アプリケーション名をプレフィックスとして設定
      schemaFilePath: "./src/resources/schema", // スキーマファイルのパスを設定
      executorCodePath: "./src/resources/executor", // 実行コードのパスを設定
      foundationModel: PARAMS.BEDROCK.FOUNDATION_MODEL_NAME, // 基盤モデルの名前を設定
      instruction: // エージェントの指示を設定
        "You are an agent that can handle various tasks related to insurance claims, including looking up claim \\ndetails, finding what paperwork is outstanding, and sending reminders. Only send reminders if you have been \\nexplicitly requested to do so. If a user asks about your functionality, provide guidance in natural language and do not include function names in the output",
      knowledgeBaseParams: {
        knowledgeBaseId, // 知識ベースIDを設定
        description: "agent knowledge base", // 知識ベースの説明を設定
      },
      schemaBucket: props.schemaBucket, // スキーマバケットを設定
    });

    // エージェントIDをCloudFormation出力として定義
    new cdk.CfnOutput(this, "AgentId", {
      value: agent.agentId, // エージェントIDを出力
      exportName: EXPORT_NAME.AGENT_ID, // エクスポート名を設定
    });
  }
}

実行コマンド例は以下のとおりです

$ cdk deploy AgentStack

実行が完了すると、AWSコンソールでAgentが作成されていることが確認できます。KnowledgeBaseとも連携していることが確認できるかと思います。

おわりに

以上が、AWS CDKを使ってPineconeとBedrockのKnowledgeBaseとAgentの環境を構築した流れとなります。
ソースコードはすべて載せられなかったので、ぜひリポジトリの方をご参照ください。

Agentを作成する環境を簡単に構築できるようになったので、ここからいろいろなAgentを作成していこうと思っています。

シンギュラリティ・ソサエティのAI Innovators Hubに参加しています。

Discussion