amplify × knowledgeBaseで簡単にRAGを構築できるようになりました！

# この記事について

先日、AWS Amplify AI kitが公開され、Bedrockのモデルを活用したAI機能を数行のコードで簡単に実装できるようになりました。
今回はawsが公開してる、AWS Amplify AI kitを使った[サンプル](https://github.com/aws-samples/amplify-ai-examples)を改変してknowledgeBaseと統合してみたので、その紹介をしたいと思います。

今回作成したリポジトリは以下です。

https://github.com/suekou/amplify-ai-examples

knowledgeBaseを使ったことがある人を対象読者として想定しています。使ったことがない人向けの記事は後日作成予定です🙏

## 基本のチャット機能
```typescript
chat: a.conversation({
    aiModel: a.ai.model("Claude 3.5 Sonnet"),
    systemPrompt: `You are a helpful assistant`,
  })
    .authorization((allow) => allow.owner()),
```
Amplifyのデータスキーマに、Conversationのroutes設定をこのように行うと、ストリーミングの対話APIが作成することができます。このconversationにはtoolを設定することができます。([ドキュメント](https://docs.amplify.aws/react/ai/conversation/tools/))。このtoolにretrieveをする関数(Query tools)を設定することで、RAGを構築します。

## RAGの構築
 1.knowledgeBaseの作成
Amplifyには、knowledgeBaseをデータスキーマやストレージのように簡潔に定義できる機能はないので、custom resourcesを使って定義します。([ドキュメント](https://docs.amplify.aws/react/build-a-backend/add-aws-services/custom-resources/))

`amplify/custom/`下に以下のようにknowledgeBaseの定義をします。
```typescript
import { appConfig } from "./../../config";
import { v4 as uuidv4 } from "uuid";
import { Construct } from "constructs";
import * as iam from "aws-cdk-lib/aws-iam";
import * as bedrock from "aws-cdk-lib/aws-bedrock";

interface KnowledgeBaseProps {
  storageArn: string;
}

export class KnowledgeBase extends Construct {
  public readonly knowledgeBaseId: string;

  constructor(scope: Construct, id: string, props: KnowledgeBaseProps) {
    super(scope, id);

    const {
      pineconeEndpoint,
      pineconeApiKeySecretArn,
      embeddingModelArn,
    } = appConfig;

    const uid = uuidv4();

    const knowledgeBaseRole = new iam.Role(this, `KnowledgeBaseRole`, {
      roleName: `knowledgeBase-role-${uid}`,
      assumedBy: new iam.ServicePrincipal("bedrock.amazonaws.com"),
      inlinePolicies: {
        inlinePolicy1: new iam.PolicyDocument({
          statements: [
            new iam.PolicyStatement({
              resources: [pineconeApiKeySecretArn],
              actions: ["secretsmanager:GetSecretValue"],
            }),
            new iam.PolicyStatement({
              resources: [embeddingModelArn],
              actions: ["bedrock:InvokeModel"],
            }),
            new iam.PolicyStatement({
              resources: [props.storageArn, `${props.storageArn}/*`],
              actions: ["s3:ListBucket", "s3:GetObject"],
            }),
          ],
        }),
      },
    });

    const knowledgeBase = new bedrock.CfnKnowledgeBase(this, "KnowledgeBase", {
      knowledgeBaseConfiguration: {
        type: "VECTOR",
        vectorKnowledgeBaseConfiguration: {
          embeddingModelArn,
        },
      },
      name: `knowledgeBase-${uid}`,
      roleArn: knowledgeBaseRole.roleArn,
      storageConfiguration: {
        pineconeConfiguration: {
          connectionString: pineconeEndpoint,
          credentialsSecretArn: pineconeApiKeySecretArn,
          fieldMapping: {
            metadataField: "metadata",
            textField: "text",
          },
        },
        type: "PINECONE",
      },
    });

    this.knowledgeBaseId = knowledgeBase.ref;

    new bedrock.CfnDataSource(this, "BedrockKnowledgeBaseDataStore", {
      name: `knowledgeBase-data-source-${uid}`,
      knowledgeBaseId: knowledgeBase.ref,
      dataSourceConfiguration: {
        s3Configuration: {
          bucketArn: props.storageArn,
          inclusionPrefixes: ["rag-resources/"],
        },
        type: "S3",
      },
    });
  }
}

```
ベクトルDBを変更したい場合は、storageConfigurationを変更してください。
これを、createStack()で既存のリソースと統合します。
```typescript
const knowledgeBaseStack = backend.createStack("knowledgeBase");
const knowledgeBase = new KnowledgeBase(knowledgeBaseStack, "KnowledgeBase", {
  storageArn: backend.storage.resources.bucket.bucketArn,
});
```

2.retrieveするためのHTTP resolverを作成する。
実際にtoolとして使用するためのHTTP resolverを`amplify/data/resolvers/kbResolver.js`に作成します。
```js
export function request(ctx) {
  const { input } = ctx.args;
  return {
    resourcePath: "/knowledgebases/<your-knowledgeBase-id>/retrieve",
    method: "POST",
    params: {
      headers: {
        "Content-Type": "application/json",
      },
      body: JSON.stringify({
        retrievalQuery: {
          text: input,
        },
      }),
    },
  };
}

export function response(ctx) {
  const responseBody = JSON.parse(ctx.result.body);
  const retrievalResults = responseBody.retrievalResults || [];

  const texts = retrievalResults.map(result => result.content.text);

  return JSON.stringify(texts);
}
```
[ドキュメント](https://docs.amplify.aws/react/ai/conversation/knowledge-base/#write-an-aws-appsync-resolver)では、`response()`が
```js
export function response(ctx) {
  return JSON.stringify(ctx.result.body);
}
```
と定義されていますが、これだとtokenの量が多くなりすぎて処理が重くなってしまったので、textのみを返しています。([レスポンスの型](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_Retrieve.html#API_agent-runtime_Retrieve_ResponseSyntax))

作成したファイルをエントリポイントとして、スキーマに加えます。
```typescript
knowledgeBase: a
    .query()
    .arguments({ input: a.string() })
    .handler(
      a.handler.custom({
        dataSource: "KnowledgeBaseDataSource",
        entry: "./resolvers/kbResolver.js",
      }),
    )
    .returns(a.string())
    .authorization((allow) => allow.authenticated()),
```
次に、これをtoolとして使えるように、定義してあるchatに加えます。
```diff_typescript
chat: a.conversation({
    aiModel: a.ai.model("Claude 3 Haiku"),
    systemPrompt: `You are a helpful assistant`,
+    tools: [
+     a.ai.dataTool({
+       name: 'searchDocumentation',
+       description: 'Performs a similarity search over the documentation',
+       query: a.ref('knowledgeBase'),
+     }),
    ]
  })
    .authorization((allow) => allow.owner()),
```

最後に`backend.ts`でHTTP resolverがretirveするための設定を行います。
```typescript

const KnowledgeBaseDataSource = backend.data.resources.graphqlApi.addHttpDataSource(
  "KnowledgeBaseDataSource",
  `https://bedrock-agent-runtime.${cdk.Stack.of(backend.data).region}.amazonaws.com`,
  {
    authorizationConfig: {
      signingRegion: cdk.Stack.of(backend.data).region,
      signingServiceName: "bedrock",
    },
  }
);

KnowledgeBaseDataSource.grantPrincipal.addToPrincipalPolicy(
  new PolicyStatement({
    resources: [
      `arn:aws:bedrock:${
        cdk.Stack.of(backend.data).region
      }:<your-aws-accountId>:knowledge-base/${knowledgeBase.knowledgeBaseId}`,
    ],
    actions: ["bedrock:Retrieve"],
  })
);
```

## 実際に動かしてみる
- 事前準備
1.pineconeのアカウントを作成し、dimensionsが1024のindexを作成する。([こちら](https://qiita.com/minorun365/items/ed680c2fc246f97985d2#1-pinecone%E3%82%A4%E3%83%B3%E3%83%87%E3%83%83%E3%82%AF%E3%82%B9%E3%81%AE%E4%BD%9C%E6%88%90%E7%84%A1%E6%96%99%E3%83%97%E3%83%A9%E3%83%B3)を参考にしてください)
2.pineconeのapi keyをシークレットマネージャーに保存する。([こちら](https://qiita.com/minorun365/items/ed680c2fc246f97985d2#2-%E3%82%B7%E3%83%BC%E3%82%AF%E3%83%AC%E3%83%83%E3%83%88%E3%83%9E%E3%83%8D%E3%83%BC%E3%82%B8%E3%83%A3%E3%83%BC%E3%81%AE%E8%A8%AD%E5%AE%9A)を参考にしてください)

- 環境構築
1.[リポジトリ](https://github.com/suekou/amplify-ai-examples)をcloneして、`npm install`もしくは`yarn`でパッケージをインストールします。
2.`claude-ai/amplify/config.ts`に事前準備で作ったpineconeのエンドポイントを`pineconeEndpoint`に、シークレットマネージャーのarnを`pineconeApiKeySecretArn`に入れます。
また、`claude-ai/amplify/backend.ts`に構築先のawsのユーザーIDを`<your-aws-accountId>`に入れます。
3.sandbox環境を立ち上げます。
```shell
cd claude-ai && npx ampx sandbox
```
4.awsのコンソールからbedrock→knowledgeBaseへ移動し、作成されたknwoledgeBaseのidを`claude-ai/amplify/data/resolvers/kbResolver.js`の`<your-knowledgeBase-id>`に入れます。
5.コンソールのknowledgeBaseからソースリンクへ行き、RAGの参照資料をアップロードします。([サポートしているファイル](https://docs.aws.amazon.com/bedrock/latest/userguide/knowledge-base-ds.html))
6.アップロードが完了したら、同期ボタンを押し、同期を完了させます。

- チャットしてみる
1. `claude-ai/`で`npm run dev`もしくは`yarn dev`を実行しlocalhostを立ち上げます。
2. 次のような画面が表示されるので、アカウントを作成します。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3522149/badffee3-da28-651d-908a-c51193cac6ef.png)
3.Create Chatボタンを押して、アップロードしたドキュメントに沿った質問をしてみましょう。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3522149/50fbc9b5-73d3-f377-4e66-50acfc203159.png)
ここでは、アップロードしてあった架空の議事録ファイルについて答えることができることが確認できました！

## リソースの削除
作ったリソースは忘れずに削除しましょう。以下のコマンドで削除できます。
```
npx ampx sandbox delete
```

## 最後に
公開されたAmplify AI kitにはまだまだ機能があるので、試したいと思います。
改善案や質問等あればコメントお願いします！