開発チームがお届けするブログリレーです!既に公開されている記事もありますので、こちらから他のメンバーの投稿もぜひチェックしてみてください!

はじめに

みなさんご存知かと思いますが、generative-ai-use-cases-jp (GenU) は aws-samples で公開されている生成 AI アプリケーションのサンプル実装です。

GenU(Generative UI)では、ユーザーのプロンプト履歴やAIの応答履歴が保存されています。しかし、これらのデータを無期限に保存すると、ストレージコストの増大などの問題が発生します。

この記事では、DynamoDBのTTL(Time to Live)機能とDynamoDB Streamsを組み合わせて、GenUのプロンプト履歴に保持期間を設定する方法を解説します。

GenUにおけるDynamoDBの活用

GenUでは、DynamoDBを使用してチャット履歴やシステムコンテキストなどのデータを保存しています。このテーブルにTTL機能の設定を行います。

// packages/cdk/lib/construct/database.ts
import { Construct } from 'constructs';
import * as ddb from 'aws-cdk-lib/aws-dynamodb';

export class Database extends Construct {
  public readonly table: ddb.Table;
  public readonly feedbackIndexName: string;
  constructor(scope: Construct, id: string) {
    super(scope, id);

    const feedbackIndexName = 'FeedbackIndex';
    const table = new ddb.Table(this, 'Table', {
      partitionKey: {
        name: 'id',
        type: ddb.AttributeType.STRING,
      },
      sortKey: {
        name: 'createdDate',
        type: ddb.AttributeType.STRING,
      },
      billingMode: ddb.BillingMode.PAY_PER_REQUEST,
      timeToLiveAttribute: 'ttl',
    });

    table.addGlobalSecondaryIndex({
      indexName: feedbackIndexName,
      partitionKey: {
        name: 'feedback',
        type: ddb.AttributeType.STRING,
      },
    });

    this.table = table;
    this.feedbackIndexName = feedbackIndexName;
  }
}

このコードでは、テーブルの作成時にtimeToLiveAttributettlに設定することで、DynamoDBのTTL機能を有効化しています。このTTL属性が設定されたレコードは、指定された時間が経過すると自動的に削除されます。

アーキテクチャの拡張

GenUのプロンプト履歴に適切な保持期間を設定するために、既存のアーキテクチャを以下のように拡張します:

  1. 既存のDynamoDBテーブルにストリーム機能を追加
  2. データ書き込みがあるたびにDynamoDB StreamsをトリガーにしたLambda関数を実行
  3. Lambda関数でTTL属性を設定

1. DynamoDBテーブルのストリーム有効化

まず、既存のDatabaseコンストラクトを拡張して、DynamoDB Streamsを有効化します。

// database.tsの拡張
import { Construct } from 'constructs';
import * as ddb from 'aws-cdk-lib/aws-dynamodb';

export class Database extends Construct {
  public readonly table: ddb.Table;
  public readonly feedbackIndexName: string;
  constructor(scope: Construct, id: string) {
    super(scope, id);

    const feedbackIndexName = 'FeedbackIndex';
    const table = new ddb.Table(this, 'Table', {
      partitionKey: {
        name: 'id',
        type: ddb.AttributeType.STRING,
      },
      sortKey: {
        name: 'createdDate',
        type: ddb.AttributeType.STRING,
      },
      billingMode: ddb.BillingMode.PAY_PER_REQUEST,
      timeToLiveAttribute: 'ttl',
      // ストリームを有効化
      stream: ddb.StreamViewType.NEW_IMAGE,
    });

    table.addGlobalSecondaryIndex({
      indexName: feedbackIndexName,
      partitionKey: {
        name: 'feedback',
        type: ddb.AttributeType.STRING,
      },
    });

    this.table = table;
    this.feedbackIndexName = feedbackIndexName;
  }
}

2. TTL設定用Lambda関数の実装

次に、DynamoDB Streamsによって起動され、TTL属性を設定するLambda関数を実装します。

// packages/cdk/lambda/setTableTTL.ts
import { DynamoDBClient } from '@aws-sdk/client-dynamodb';
import {
  DynamoDBDocumentClient,
  UpdateCommand,
} from '@aws-sdk/lib-dynamodb';
import { DynamoDBStreamEvent, Context } from 'aws-lambda';

const TABLE_NAME: string = process.env.TABLE_NAME!;
const TTL_DAYS: number = parseInt(process.env.TTL_DAYS || '30', 10);
const dynamoDb = new DynamoDBClient({});
const dynamoDbDocument = DynamoDBDocumentClient.from(dynamoDb);

export const handler = async (event: DynamoDBStreamEvent, context: Context): Promise<void> => {
  console.log('TTL設定関数が呼び出されました');

  // バッチ処理のためのPromiseを格納する配列
  const promises = [];

  for (const record of event.Records) {
    // INSERT(新規作成)イベントのみを処理
    if (record.eventName !== 'INSERT') {
      continue;
    }

    // 新しく作成されたレコードのデータを取得
    const newImage = record.dynamodb?.NewImage;
    if (!newImage) {
      continue;
    }

    // DynamoDBのイメージをJavaScriptオブジェクトに変換
    const item = AWS.DynamoDB.Converter.unmarshall(newImage);

    // TTLが既に設定されている場合はスキップ
    if (item.ttl) {
      continue;
    }

    // 現在時刻からTTL期間(秒)を計算
    const now = Math.floor(Date.now() / 1000);
    const ttlValue = now + (TTL_DAYS * 24 * 60 * 60);

    console.log(`レコード ${item.id}:${item.createdDate} にTTL=${ttlValue}を設定します`);

    // レコードを更新してTTL属性を追加
    const params = {
      TableName: TABLE_NAME,
      Key: {
        id: item.id,
        createdDate: item.createdDate,
      },
      UpdateExpression: 'SET ttl = :ttl',
      ExpressionAttributeValues: {
        ':ttl': ttlValue,
      },
    };

    promises.push(dynamoDbDocument.send(new UpdateCommand(params)));
  }

  // すべての更新を並行して実行
  await Promise.all(promises);
  console.log(`${promises.length}件のレコードにTTLを設定しました`);
};

このLambda関数は、テーブルに挿入される全てのレコードを対象にTTL属性を設定します。プレフィックスやデータタイプによるフィルタリングはせず、すべてのデータに対して保持期間を設定します。

3. CDKスタックにLambda関数を追加

最後に、この関数をGenUのCDKスタックに統合します。既存のgenerative-ai-use-cases-stack.tsを拡張します。

// packages/cdk/lib/generative-ai-use-cases-stack.tsの拡張
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as lambdaEventSources from 'aws-cdk-lib/aws-lambda-eventSources';

// 既存のコードの続き
// TTL設定用Lambda関数の作成
const ttlSetterFunction = new lambda.NodejsFunction(this, 'TTLSetterFunction', {
  runtime: lambda.Runtime.NODEJS_LATEST,
  entry: './lambda/setTableTTL.ts',
  timeout: cdk.Duration.minutes(5),
  environment: {
    TABLE_NAME: database.table.tableName,
    TTL_DAYS: '30', // 保持期間を30日に設定
  },
});

// DynamoDB StreamsをLambda関数のトリガーとして設定
ttlSetterFunction.addEventSource(
  new lambdaEventSources.DynamoEventSource(database.table, {
    startingPosition: lambda.StartingPosition.LATEST,
    batchSize: 100,
    retryAttempts: 3,
  })
);

// Lambda関数にDynamoDBテーブルへの書き込み権限を付与
database.table.grantWriteData(ttlSetterFunction);

GenUアプリケーションでの実装の影響

この実装では、既存のGenUアプリケーションコードには変更は不要です。既存のrepository.tsなどのデータアクセス層は、そのまま利用できます。テーブルに保存されるすべてのデータに対してTTLが設定され、設定された期間後に自動的に削除されます。

これにより、アプリケーション側の変更リスクを最小限に抑えながら、データの保持期間管理を実現できます。すべてのデータタイプ(チャット履歴、プロンプト、システムコンテキストなど)に対して一貫した保持ポリシーが適用されます。

メリットと注意点

メリット

  1. データライフサイクル管理の自動化: TTLによって全てのデータが自動的に削除されるため、手動でのクリーンアップが不要
  2. コスト削減: 不要なデータを削除することでDynamoDBのストレージコストを削減
  3. プライバシー保護: ユーザーデータを必要以上に長く保持しない
  4. パフォーマンス向上: テーブルサイズが小さく保たれるため、クエリのパフォーマンスが向上
  5. 既存コードの活用: GenUの既存のインフラストラクチャとコードを拡張して実装できる
  6. アプリケーションコードへの影響なし: データアクセス層に変更を加えずに実装できる
  7. 一貫したデータポリシー: 全てのデータタイプに対して同じ保持期間ポリシーを適用

注意点

  1. TTL削除のタイミング: DynamoDBのTTL削除は保証された時間に必ず行われるわけではなく、通常は48時間以内に行われます
  2. 課金への影響: TTLによる削除はWCU(書き込みキャパシティユニット)を消費しませんが、ストリームは削除イベントを生成します
  3. ストリーム処理のエラーハンドリング: Lambda関数がエラーになった場合のリトライ戦略を考慮する必要があります
  4. ストリームレコードの期限: DynamoDB Streamsのレコードは24時間後に期限切れになるため、長時間のサービス停止に注意が必要です
  5. データタイプによる区別がない: すべてのレコードに同じTTL期間が適用されるため、データタイプによって保持期間を変えたい場合は追加のロジックが必要

まとめ

DynamoDBのTTL機能とDynamoDB Streamsを組み合わせてGenUのデータの保持期間設定を実装しました。テーブル内の全レコードに対してTTLを設定することで、上記のメリットが得られるのはとてもありがたいですね!

参考にしていただけたら幸いです!