aws-genai-llm-chatbot

1
import * as path from "path";
2
import * as cdk from "aws-cdk-lib";
3
import { Construct } from "constructs";
4
import { SystemConfig } from "../../shared/types";
5
import { Shared } from "../../shared";
6
import { RagDynamoDBTables } from "../rag-dynamodb-tables";
7
import * as sfn from "aws-cdk-lib/aws-stepfunctions";
8
import * as tasks from "aws-cdk-lib/aws-stepfunctions-tasks";
9
import * as lambda from "aws-cdk-lib/aws-lambda";
10
import * as logs from "aws-cdk-lib/aws-logs";
11
import * as rds from "aws-cdk-lib/aws-rds";
12
import { RemovalPolicy } from "aws-cdk-lib";
13

14
export interface CreateAuroraWorkspaceProps {
15
  readonly config: SystemConfig;
16
  readonly shared: Shared;
17
  readonly ragDynamoDBTables: RagDynamoDBTables;
18
  readonly dbCluster: rds.DatabaseCluster;
19
}
20

21
export class CreateAuroraWorkspace extends Construct {
22
  public readonly stateMachine: sfn.StateMachine;
23

24
  constructor(scope: Construct, id: string, props: CreateAuroraWorkspaceProps) {
25
    super(scope, id);
26

27
    const createFunction = new lambda.Function(
28
      this,
29
      "CreateAuroraWorkspaceFunction",
30
      {
31
        vpc: props.shared.vpc,
32
        code: props.shared.sharedCode.bundleWithLambdaAsset(
33
          path.join(__dirname, "./functions/create-workflow/create")
34
        ),
35
        runtime: props.shared.pythonRuntime,
36
        architecture: props.shared.lambdaArchitecture,
37
        handler: "index.lambda_handler",
38
        layers: [props.shared.powerToolsLayer, props.shared.commonLayer],
39
        timeout: cdk.Duration.minutes(5),
40
        logRetention: logs.RetentionDays.ONE_WEEK,
41
        environment: {
42
          ...props.shared.defaultEnvironmentVariables,
43
          AURORA_DB_SECRET_ID: props.dbCluster.secret?.secretArn as string,
44
          WORKSPACES_TABLE_NAME:
45
            props.ragDynamoDBTables.workspacesTable.tableName,
46
          WORKSPACES_BY_OBJECT_TYPE_INDEX_NAME:
47
            props.ragDynamoDBTables.workspacesByObjectTypeIndexName,
48
        },
49
      }
50
    );
51

52
    props.dbCluster.secret?.grantRead(createFunction);
53
    props.dbCluster.connections.allowDefaultPortFrom(createFunction);
54
    props.ragDynamoDBTables.workspacesTable.grantReadWriteData(createFunction);
55

56
    const handleError = new tasks.DynamoUpdateItem(this, "HandleError", {
57
      table: props.ragDynamoDBTables.workspacesTable,
58
      key: {
59
        workspace_id: tasks.DynamoAttributeValue.fromString(
60
          sfn.JsonPath.stringAt("$.workspace_id")
61
        ),
62
        object_type: tasks.DynamoAttributeValue.fromString("workspace"),
63
      },
64
      updateExpression: "set #status = :error",
65
      expressionAttributeNames: {
66
        "#status": "status",
67
      },
68
      expressionAttributeValues: {
69
        ":error": tasks.DynamoAttributeValue.fromString("error"),
70
      },
71
    }).next(
72
      new sfn.Fail(this, "Fail", {
73
        cause: "Workspace creation failed",
74
      })
75
    );
76

77
    const setCreating = new tasks.DynamoUpdateItem(this, "SetCreating", {
78
      table: props.ragDynamoDBTables.workspacesTable,
79
      key: {
80
        workspace_id: tasks.DynamoAttributeValue.fromString(
81
          sfn.JsonPath.stringAt("$.workspace_id")
82
        ),
83
        object_type: tasks.DynamoAttributeValue.fromString("workspace"),
84
      },
85
      updateExpression: "set #status=:statusValue",
86
      expressionAttributeNames: {
87
        "#status": "status",
88
      },
89
      expressionAttributeValues: {
90
        ":statusValue": tasks.DynamoAttributeValue.fromString("creating"),
91
      },
92
      resultPath: sfn.JsonPath.DISCARD,
93
    });
94

95
    const setReady = new tasks.DynamoUpdateItem(this, "SetReady", {
96
      table: props.ragDynamoDBTables.workspacesTable,
97
      key: {
98
        workspace_id: tasks.DynamoAttributeValue.fromString(
99
          sfn.JsonPath.stringAt("$.workspace_id")
100
        ),
101
        object_type: tasks.DynamoAttributeValue.fromString("workspace"),
102
      },
103
      updateExpression: "set #status=:statusValue",
104
      expressionAttributeNames: {
105
        "#status": "status",
106
      },
107
      expressionAttributeValues: {
108
        ":statusValue": tasks.DynamoAttributeValue.fromString("ready"),
109
      },
110
      resultPath: sfn.JsonPath.DISCARD,
111
    });
112

113
    const createTask = new tasks.LambdaInvoke(this, "Create", {
114
      lambdaFunction: createFunction,
115
      resultPath: "$.createResult",
116
    }).addCatch(handleError, {
117
      errors: ["States.ALL"],
118
      resultPath: "$.createResult",
119
    });
120

121
    const workflow = setCreating
122
      .next(createTask)
123
      .next(setReady)
124
      .next(new sfn.Succeed(this, "Success"));
125

126
    const logGroup = new logs.LogGroup(
127
      this,
128
      "CreateAuroraWorkspaceSMLogGroup",
129
      {
130
        removalPolicy: RemovalPolicy.DESTROY,
131
      }
132
    );
133

134
    const stateMachine = new sfn.StateMachine(this, "CreateAuroraWorkspace", {
135
      definitionBody: sfn.DefinitionBody.fromChainable(workflow),
136
      timeout: cdk.Duration.minutes(5),
137
      comment: "Create Aurora Workspace Workflow",
138
      tracingEnabled: true,
139
      logs: {
140
        destination: logGroup,
141
        level: sfn.LogLevel.ALL,
142
      },
143
    });
144

145
    this.stateMachine = stateMachine;
146
  }
147
}
148