aws-genai-llm-chatbot

1
import * as path from "path";
2
import * as cdk from "aws-cdk-lib";
3
import { SageMakerModelEndpoint, SystemConfig } from "../shared/types";
4
import { Construct } from "constructs";
5
import { RagEngines } from "../rag-engines";
6
import * as cognito from "aws-cdk-lib/aws-cognito";
7
import * as dynamodb from "aws-cdk-lib/aws-dynamodb";
8
import * as ec2 from "aws-cdk-lib/aws-ec2";
9
import * as iam from "aws-cdk-lib/aws-iam";
10
import * as lambda from "aws-cdk-lib/aws-lambda";
11
import * as logs from "aws-cdk-lib/aws-logs";
12
import * as ssm from "aws-cdk-lib/aws-ssm";
13
import { Shared } from "../shared";
14
import * as appsync from "aws-cdk-lib/aws-appsync";
15
import { parse } from "graphql";
16
import { readFileSync } from "fs";
17
import * as s3 from "aws-cdk-lib/aws-s3";
18

19
export interface ApiResolversProps {
20
  readonly shared: Shared;
21
  readonly config: SystemConfig;
22
  readonly ragEngines?: RagEngines;
23
  readonly userPool: cognito.UserPool;
24
  readonly sessionsTable: dynamodb.Table;
25
  readonly byUserIdIndex: string;
26
  readonly userFeedbackBucket: s3.Bucket;
27
  readonly modelsParameter: ssm.StringParameter;
28
  readonly models: SageMakerModelEndpoint[];
29
  readonly api: appsync.GraphqlApi;
30
}
31

32
export class ApiResolvers extends Construct {
33
  constructor(scope: Construct, id: string, props: ApiResolversProps) {
34
    super(scope, id);
35

36
    const apiSecurityGroup = new ec2.SecurityGroup(this, "ApiSecurityGroup", {
37
      vpc: props.shared.vpc,
38
    });
39

40
    const appSyncLambdaResolver = new lambda.Function(
41
      this,
42
      "GraphQLApiHandler",
43
      {
44
        code: props.shared.sharedCode.bundleWithLambdaAsset(
45
          path.join(__dirname, "./functions/api-handler")
46
        ),
47
        handler: "index.handler",
48
        runtime: props.shared.pythonRuntime,
49
        architecture: props.shared.lambdaArchitecture,
50
        timeout: cdk.Duration.minutes(10),
51
        memorySize: 512,
52
        tracing: lambda.Tracing.ACTIVE,
53
        logRetention: logs.RetentionDays.ONE_WEEK,
54
        layers: [props.shared.powerToolsLayer, props.shared.commonLayer],
55
        vpc: props.shared.vpc,
56
        securityGroups: [apiSecurityGroup],
57
        vpcSubnets: props.shared.vpc.privateSubnets as ec2.SubnetSelection,
58
        environment: {
59
          ...props.shared.defaultEnvironmentVariables,
60
          CONFIG_PARAMETER_NAME: props.shared.configParameter.parameterName,
61
          MODELS_PARAMETER_NAME: props.modelsParameter.parameterName,
62
          X_ORIGIN_VERIFY_SECRET_ARN:
63
            props.shared.xOriginVerifySecret.secretArn,
64
          API_KEYS_SECRETS_ARN: props.shared.apiKeysSecret.secretArn,
65
          SESSIONS_TABLE_NAME: props.sessionsTable.tableName,
66
          SESSIONS_BY_USER_ID_INDEX_NAME: props.byUserIdIndex,
67
          USER_FEEDBACK_BUCKET_NAME: props.userFeedbackBucket?.bucketName ?? "",
68
          UPLOAD_BUCKET_NAME: props.ragEngines?.uploadBucket?.bucketName ?? "",
69
          PROCESSING_BUCKET_NAME:
70
            props.ragEngines?.processingBucket?.bucketName ?? "",
71
          AURORA_DB_SECRET_ID: props.ragEngines?.auroraPgVector?.database
72
            ?.secret?.secretArn as string,
73
          WORKSPACES_TABLE_NAME:
74
            props.ragEngines?.workspacesTable.tableName ?? "",
75
          WORKSPACES_BY_OBJECT_TYPE_INDEX_NAME:
76
            props.ragEngines?.workspacesByObjectTypeIndexName ?? "",
77
          DOCUMENTS_TABLE_NAME:
78
            props.ragEngines?.documentsTable.tableName ?? "",
79
          DOCUMENTS_BY_COMPOUND_KEY_INDEX_NAME:
80
            props.ragEngines?.documentsByCompountKeyIndexName ?? "",
81
          DOCUMENTS_BY_STATUS_INDEX:
82
            props.ragEngines?.documentsByStatusIndexName ?? "",
83
          SAGEMAKER_RAG_MODELS_ENDPOINT:
84
            props.ragEngines?.sageMakerRagModels?.model.endpoint
85
              ?.attrEndpointName ?? "",
86
          DELETE_WORKSPACE_WORKFLOW_ARN:
87
            props.ragEngines?.deleteWorkspaceWorkflow?.stateMachineArn ?? "",
88
          CREATE_AURORA_WORKSPACE_WORKFLOW_ARN:
89
            props.ragEngines?.auroraPgVector?.createAuroraWorkspaceWorkflow
90
              ?.stateMachineArn ?? "",
91
          CREATE_OPEN_SEARCH_WORKSPACE_WORKFLOW_ARN:
92
            props.ragEngines?.openSearchVector
93
              ?.createOpenSearchWorkspaceWorkflow?.stateMachineArn ?? "",
94
          CREATE_KENDRA_WORKSPACE_WORKFLOW_ARN:
95
            props.ragEngines?.kendraRetrieval?.createKendraWorkspaceWorkflow
96
              ?.stateMachineArn ?? "",
97
          FILE_IMPORT_WORKFLOW_ARN:
98
            props.ragEngines?.fileImportWorkflow?.stateMachineArn ?? "",
99
          WEBSITE_CRAWLING_WORKFLOW_ARN:
100
            props.ragEngines?.websiteCrawlingWorkflow?.stateMachineArn ?? "",
101
          OPEN_SEARCH_COLLECTION_ENDPOINT:
102
            props.ragEngines?.openSearchVector?.openSearchCollectionEndpoint ??
103
            "",
104
          DEFAULT_KENDRA_INDEX_ID:
105
            props.ragEngines?.kendraRetrieval?.kendraIndex?.attrId ?? "",
106
          DEFAULT_KENDRA_INDEX_NAME:
107
            props.ragEngines?.kendraRetrieval?.kendraIndex?.name ?? "",
108
          DEFAULT_KENDRA_S3_DATA_SOURCE_ID:
109
            props.ragEngines?.kendraRetrieval?.kendraS3DataSource?.attrId ?? "",
110
          DEFAULT_KENDRA_S3_DATA_SOURCE_BUCKET_NAME:
111
            props.ragEngines?.kendraRetrieval?.kendraS3DataSourceBucket
112
              ?.bucketName ?? "",
113
          RSS_FEED_INGESTOR_FUNCTION:
114
            props.ragEngines?.dataImport.rssIngestorFunction?.functionArn ?? "",
115
        },
116
      }
117
    );
118

119
    function addPermissions(apiHandler: lambda.Function) {
120
      if (props.ragEngines?.workspacesTable) {
121
        props.ragEngines.workspacesTable.grantReadWriteData(apiHandler);
122
      }
123

124
      if (props.ragEngines?.documentsTable) {
125
        props.ragEngines.documentsTable.grantReadWriteData(apiHandler);
126
        props.ragEngines?.dataImport.rssIngestorFunction?.grantInvoke(
127
          apiHandler
128
        );
129
      }
130

131
      if (props.ragEngines?.auroraPgVector) {
132
        props.ragEngines.auroraPgVector.database.secret?.grantRead(apiHandler);
133
        props.ragEngines.auroraPgVector.database.connections.allowDefaultPortFrom(
134
          apiHandler
135
        );
136

137
        props.ragEngines.auroraPgVector.createAuroraWorkspaceWorkflow.grantStartExecution(
138
          apiHandler
139
        );
140
      }
141

142
      if (props.ragEngines?.openSearchVector) {
143
        apiHandler.addToRolePolicy(
144
          new iam.PolicyStatement({
145
            actions: ["aoss:APIAccessAll"],
146
            resources: [
147
              props.ragEngines?.openSearchVector.openSearchCollection.attrArn,
148
            ],
149
          })
150
        );
151

152
        props.ragEngines.openSearchVector.createOpenSearchWorkspaceWorkflow.grantStartExecution(
153
          apiHandler
154
        );
155
      }
156

157
      if (props.ragEngines?.kendraRetrieval) {
158
        props.ragEngines.kendraRetrieval.createKendraWorkspaceWorkflow.grantStartExecution(
159
          apiHandler
160
        );
161

162
        props.ragEngines?.kendraRetrieval?.kendraS3DataSourceBucket?.grantReadWrite(
163
          apiHandler
164
        );
165

166
        if (props.ragEngines.kendraRetrieval.kendraIndex) {
167
          apiHandler.addToRolePolicy(
168
            new iam.PolicyStatement({
169
              actions: [
170
                "kendra:Retrieve",
171
                "kendra:Query",
172
                "kendra:BatchDeleteDocument",
173
                "kendra:BatchPutDocument",
174
                "kendra:StartDataSourceSyncJob",
175
                "kendra:DescribeDataSourceSyncJob",
176
                "kendra:StopDataSourceSyncJob",
177
                "kendra:ListDataSourceSyncJobs",
178
                "kendra:ListDataSources",
179
                "kendra:DescribeIndex",
180
              ],
181
              resources: [
182
                props.ragEngines.kendraRetrieval.kendraIndex.attrArn,
183
                `${props.ragEngines.kendraRetrieval.kendraIndex.attrArn}/*`,
184
              ],
185
            })
186
          );
187
        }
188

189
        for (const item of props.config.rag.engines.kendra.external ?? []) {
190
          if (item.roleArn) {
191
            apiHandler.addToRolePolicy(
192
              new iam.PolicyStatement({
193
                actions: ["sts:AssumeRole"],
194
                resources: [item.roleArn],
195
              })
196
            );
197
          } else {
198
            apiHandler.addToRolePolicy(
199
              new iam.PolicyStatement({
200
                actions: ["kendra:Retrieve", "kendra:Query"],
201
                resources: [
202
                  `arn:${cdk.Aws.PARTITION}:kendra:${
203
                    item.region ?? cdk.Aws.REGION
204
                  }:${cdk.Aws.ACCOUNT_ID}:index/${item.kendraId}`,
205
                ],
206
              })
207
            );
208
          }
209
        }
210
      }
211

212
      if (props.ragEngines?.fileImportWorkflow) {
213
        props.ragEngines.fileImportWorkflow.grantStartExecution(apiHandler);
214
      }
215

216
      if (props.ragEngines?.websiteCrawlingWorkflow) {
217
        props.ragEngines.websiteCrawlingWorkflow.grantStartExecution(
218
          apiHandler
219
        );
220
      }
221

222
      if (props.ragEngines?.deleteWorkspaceWorkflow) {
223
        props.ragEngines.deleteWorkspaceWorkflow.grantStartExecution(
224
          apiHandler
225
        );
226
      }
227

228
      if (props.ragEngines?.sageMakerRagModels) {
229
        apiHandler.addToRolePolicy(
230
          new iam.PolicyStatement({
231
            actions: ["sagemaker:InvokeEndpoint"],
232
            resources: [props.ragEngines.sageMakerRagModels.model.endpoint.ref],
233
          })
234
        );
235
      }
236

237
      for (const model of props.models) {
238
        apiHandler.addToRolePolicy(
239
          new iam.PolicyStatement({
240
            actions: ["sagemaker:InvokeEndpoint"],
241
            resources: [model.endpoint.ref],
242
          })
243
        );
244
      }
245

246
      apiHandler.addToRolePolicy(
247
        new iam.PolicyStatement({
248
          actions: [
249
            "comprehend:DetectDominantLanguage",
250
            "comprehend:DetectSentiment",
251
          ],
252
          resources: ["*"],
253
        })
254
      );
255

256
      props.shared.xOriginVerifySecret.grantRead(apiHandler);
257
      props.shared.apiKeysSecret.grantRead(apiHandler);
258
      props.shared.configParameter.grantRead(apiHandler);
259
      props.modelsParameter.grantRead(apiHandler);
260
      props.sessionsTable.grantReadWriteData(apiHandler);
261
      props.userFeedbackBucket.grantReadWrite(apiHandler);
262
      props.ragEngines?.uploadBucket.grantReadWrite(apiHandler);
263
      props.ragEngines?.processingBucket.grantReadWrite(apiHandler);
264

265
      if (props.config.bedrock?.enabled) {
266
        apiHandler.addToRolePolicy(
267
          new iam.PolicyStatement({
268
            actions: [
269
              "bedrock:ListFoundationModels",
270
              "bedrock:ListCustomModels",
271
              "bedrock:InvokeModel",
272
              "bedrock:InvokeModelWithResponseStream",
273
            ],
274
            resources: ["*"],
275
          })
276
        );
277

278
        if (props.config.bedrock?.roleArn) {
279
          apiHandler.addToRolePolicy(
280
            new iam.PolicyStatement({
281
              actions: ["sts:AssumeRole"],
282
              resources: [props.config.bedrock.roleArn],
283
            })
284
          );
285
        }
286
      }
287
    }
288

289
    addPermissions(appSyncLambdaResolver);
290

291
    props.ragEngines?.openSearchVector?.addToAccessPolicy(
292
      "graphql-api",
293
      [appSyncLambdaResolver.role?.roleArn],
294
      ["aoss:DescribeIndex", "aoss:ReadDocument", "aoss:WriteDocument"]
295
    );
296

297
    const functionDataSource = props.api.addLambdaDataSource(
298
      "proxyResolverFunction",
299
      appSyncLambdaResolver
300
    );
301

302
    const schema = parse(
303
      readFileSync("lib/chatbot-api/schema/schema.graphql", "utf8")
304
    );
305

306
    function addResolvers(operationType: string) {
307
      const fieldNames = (
308
        schema.definitions
309
          .filter((x) => x.kind == "ObjectTypeDefinition")
310
          .filter((y: any) => y.name.value == operationType)[0] as any
311
      ).fields.map((z: any) => z.name.value);
312

313
      for (const fieldName of fieldNames) {
314
        // These resolvers are added by the Realtime API
315
        if (fieldName == "sendQuery" || fieldName == "publishResponse") {
316
          continue;
317
        }
318
        props.api.createResolver(`${fieldName}-resolver`, {
319
          typeName: operationType,
320
          fieldName: fieldName,
321
          dataSource: functionDataSource,
322
        });
323
      }
324
    }
325

326
    addResolvers("Query");
327
    addResolvers("Mutation");
328
  }
329
}
330