aws-genai-llm-chatbot

1
import * as cdk from "aws-cdk-lib";
2
import * as iam from "aws-cdk-lib/aws-iam";
3
import * as lambda from "aws-cdk-lib/aws-lambda";
4
import * as logs from "aws-cdk-lib/aws-logs";
5
import * as sfn from "aws-cdk-lib/aws-stepfunctions";
6
import * as tasks from "aws-cdk-lib/aws-stepfunctions-tasks";
7
import { Construct } from "constructs";
8
import * as path from "path";
9
import { Shared } from "../../shared";
10
import { SystemConfig } from "../../shared/types";
11
import { AuroraPgVector } from "../aurora-pgvector";
12
import { DataImport } from "../data-import";
13
import { KendraRetrieval } from "../kendra-retrieval";
14
import { OpenSearchVector } from "../opensearch-vector";
15
import { RagDynamoDBTables } from "../rag-dynamodb-tables";
16
import { RemovalPolicy } from "aws-cdk-lib";
17

18
export interface DeleteWorkspaceProps {
19
  readonly config: SystemConfig;
20
  readonly shared: Shared;
21
  readonly dataImport: DataImport;
22
  readonly ragDynamoDBTables: RagDynamoDBTables;
23
  readonly auroraPgVector?: AuroraPgVector;
24
  readonly openSearchVector?: OpenSearchVector;
25
  readonly kendraRetrieval?: KendraRetrieval;
26
}
27

28
export class DeleteWorkspace extends Construct {
29
  public readonly stateMachine?: sfn.StateMachine;
30

31
  constructor(scope: Construct, id: string, props: DeleteWorkspaceProps) {
32
    super(scope, id);
33

34
    const deleteFunction = new lambda.Function(
35
      this,
36
      "DeleteWorkspaceFunction",
37
      {
38
        vpc: props.shared.vpc,
39
        code: props.shared.sharedCode.bundleWithLambdaAsset(
40
          path.join(__dirname, "./functions/delete-workspace-workflow/delete")
41
        ),
42
        runtime: props.shared.pythonRuntime,
43
        architecture: props.shared.lambdaArchitecture,
44
        handler: "index.lambda_handler",
45
        layers: [props.shared.powerToolsLayer, props.shared.commonLayer],
46
        timeout: cdk.Duration.minutes(15),
47
        logRetention: logs.RetentionDays.ONE_WEEK,
48
        environment: {
49
          ...props.shared.defaultEnvironmentVariables,
50
          AURORA_DB_SECRET_ID: props.auroraPgVector?.database.secret
51
            ?.secretArn as string,
52
          UPLOAD_BUCKET_NAME: props.dataImport.uploadBucket.bucketName,
53
          PROCESSING_BUCKET_NAME: props.dataImport.processingBucket.bucketName,
54
          WORKSPACES_TABLE_NAME:
55
            props.ragDynamoDBTables.workspacesTable.tableName,
56
          WORKSPACES_BY_OBJECT_TYPE_INDEX_NAME:
57
            props.ragDynamoDBTables.workspacesByObjectTypeIndexName,
58
          DOCUMENTS_TABLE_NAME:
59
            props.ragDynamoDBTables?.documentsTable.tableName ?? "",
60
          DOCUMENTS_BY_COMPOUND_KEY_INDEX_NAME:
61
            props.ragDynamoDBTables?.documentsByCompoundKeyIndexName ?? "",
62
          DOCUMENTS_BY_STATUS_INDEX:
63
            props.ragDynamoDBTables.documentsByStatusIndexName ?? "",
64
          DEFAULT_KENDRA_S3_DATA_SOURCE_BUCKET_NAME:
65
            props.kendraRetrieval?.kendraS3DataSourceBucket?.bucketName ?? "",
66
          OPEN_SEARCH_COLLECTION_ENDPOINT:
67
            props.openSearchVector?.openSearchCollectionEndpoint ?? "",
68
        },
69
      }
70
    );
71

72
    if (props.auroraPgVector) {
73
      props.auroraPgVector.database.secret?.grantRead(deleteFunction);
74
      props.auroraPgVector.database.connections.allowDefaultPortFrom(
75
        deleteFunction
76
      );
77
    }
78

79
    if (props.openSearchVector) {
80
      deleteFunction.addToRolePolicy(
81
        new iam.PolicyStatement({
82
          actions: [
83
            "aoss:APIAccessAll",
84
            "aoss:DescribeIndex",
85
            "aoss:DeleteIndex",
86
          ],
87
          resources: [props.openSearchVector.openSearchCollection.attrArn],
88
        })
89
      );
90

91
      props.openSearchVector.addToAccessPolicy(
92
        "delete-workspace",
93
        [deleteFunction.role?.roleArn],
94
        [
95
          "aoss:DeleteIndex",
96
          "aoss:DescribeIndex",
97
          "aoss:ReadDocument",
98
          "aoss:WriteDocument",
99
        ]
100
      );
101
    }
102

103
    props.dataImport.uploadBucket.grantReadWrite(deleteFunction);
104
    props.dataImport.processingBucket.grantReadWrite(deleteFunction);
105
    props.kendraRetrieval?.kendraS3DataSourceBucket?.grantReadWrite(
106
      deleteFunction
107
    );
108
    props.ragDynamoDBTables.workspacesTable.grantReadWriteData(deleteFunction);
109
    props.ragDynamoDBTables.documentsTable.grantReadWriteData(deleteFunction);
110

111
    const handleError = new tasks.DynamoUpdateItem(this, "HandleError", {
112
      table: props.ragDynamoDBTables.workspacesTable,
113
      key: {
114
        workspace_id: tasks.DynamoAttributeValue.fromString(
115
          sfn.JsonPath.stringAt("$.workspace_id")
116
        ),
117
        object_type: tasks.DynamoAttributeValue.fromString("workspace"),
118
      },
119
      updateExpression: "set #status = :error",
120
      expressionAttributeNames: {
121
        "#status": "status",
122
      },
123
      expressionAttributeValues: {
124
        ":error": tasks.DynamoAttributeValue.fromString("error"),
125
      },
126
    }).next(
127
      new sfn.Fail(this, "Fail", {
128
        cause: "Workspace deletion failed",
129
      })
130
    );
131

132
    const setDeleting = new tasks.DynamoUpdateItem(this, "SetDeleting", {
133
      table: props.ragDynamoDBTables.workspacesTable,
134
      key: {
135
        workspace_id: tasks.DynamoAttributeValue.fromString(
136
          sfn.JsonPath.stringAt("$.workspace_id")
137
        ),
138
        object_type: tasks.DynamoAttributeValue.fromString("workspace"),
139
      },
140
      updateExpression: "set #status=:statusValue",
141
      expressionAttributeNames: {
142
        "#status": "status",
143
      },
144
      expressionAttributeValues: {
145
        ":statusValue": tasks.DynamoAttributeValue.fromString("deleting"),
146
      },
147
      resultPath: sfn.JsonPath.DISCARD,
148
    });
149

150
    const deleteTask = new tasks.LambdaInvoke(this, "Delete", {
151
      lambdaFunction: deleteFunction,
152
      resultPath: "$.deleteResult",
153
    }).addCatch(handleError, {
154
      errors: ["States.ALL"],
155
      resultPath: "$.deleteResult",
156
    });
157

158
    const workflow = setDeleting
159
      .next(deleteTask)
160
      .next(new sfn.Succeed(this, "Success"));
161

162
    const logGroup = new logs.LogGroup(this, "DeleteWorkspaceSMLogGroup", {
163
      removalPolicy: RemovalPolicy.DESTROY,
164
    });
165

166
    const stateMachine = new sfn.StateMachine(this, "DeleteWorkspace", {
167
      definitionBody: sfn.DefinitionBody.fromChainable(workflow),
168
      timeout: cdk.Duration.minutes(5),
169
      comment: "Delete Workspace Workflow",
170
      tracingEnabled: true,
171
      logs: {
172
        destination: logGroup,
173
        level: sfn.LogLevel.ALL,
174
      },
175
    });
176

177
    this.stateMachine = stateMachine;
178
  }
179
}
180