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 oss from "aws-cdk-lib/aws-opensearchserverless";
6
import * as sfn from "aws-cdk-lib/aws-stepfunctions";
7
import * as tasks from "aws-cdk-lib/aws-stepfunctions-tasks";
8
import { Construct } from "constructs";
9
import * as path from "path";
10
import { Shared } from "../../shared";
11
import { SystemConfig } from "../../shared/types";
12
import { RagDynamoDBTables } from "../rag-dynamodb-tables";
13
import { RemovalPolicy } from "aws-cdk-lib";
14

15
export interface CreateOpenSearchWorkspaceProps {
16
  readonly config: SystemConfig;
17
  readonly shared: Shared;
18
  readonly ragDynamoDBTables: RagDynamoDBTables;
19
  readonly openSearchCollectionName: string;
20
  readonly openSearchCollection: oss.CfnCollection;
21
  readonly collectionEndpoint: string;
22
}
23

24
export class CreateOpenSearchWorkspace extends Construct {
25
  public readonly stateMachine: sfn.StateMachine;
26
  public readonly createWorkspaceRole?: iam.IRole;
27

28
  constructor(
29
    scope: Construct,
30
    id: string,
31
    props: CreateOpenSearchWorkspaceProps
32
  ) {
33
    super(scope, id);
34

35
    const createFunction = new lambda.Function(
36
      this,
37
      "CreateOpenSearchWorkspaceFunction",
38
      {
39
        vpc: props.shared.vpc,
40
        code: props.shared.sharedCode.bundleWithLambdaAsset(
41
          path.join(__dirname, "./functions/create-workflow/create")
42
        ),
43
        runtime: props.shared.pythonRuntime,
44
        architecture: props.shared.lambdaArchitecture,
45
        handler: "index.lambda_handler",
46
        layers: [props.shared.powerToolsLayer, props.shared.commonLayer],
47
        timeout: cdk.Duration.minutes(5),
48
        logRetention: logs.RetentionDays.ONE_WEEK,
49
        environment: {
50
          ...props.shared.defaultEnvironmentVariables,
51
          WORKSPACES_TABLE_NAME:
52
            props.ragDynamoDBTables.workspacesTable.tableName,
53
          WORKSPACES_BY_OBJECT_TYPE_INDEX_NAME:
54
            props.ragDynamoDBTables.workspacesByObjectTypeIndexName,
55
          OPEN_SEARCH_COLLECTION_NAME: props.openSearchCollectionName,
56
          OPEN_SEARCH_COLLECTION_ENDPOINT: props.collectionEndpoint,
57
          OPEN_SEARCH_COLLECTION_ENDPOINT_PORT: "443",
58
        },
59
      }
60
    );
61

62
    props.ragDynamoDBTables.workspacesTable.grantReadWriteData(createFunction);
63
    createFunction.addToRolePolicy(
64
      new iam.PolicyStatement({
65
        actions: [
66
          "aoss:APIAccessAll",
67
          "aoss:DescribeIndex",
68
          "aoss:CreateIndex",
69
        ],
70
        resources: [props.openSearchCollection.attrArn],
71
      })
72
    );
73

74
    const handleError = new tasks.DynamoUpdateItem(this, "HandleError", {
75
      table: props.ragDynamoDBTables.workspacesTable,
76
      key: {
77
        workspace_id: tasks.DynamoAttributeValue.fromString(
78
          sfn.JsonPath.stringAt("$.workspace_id")
79
        ),
80
        object_type: tasks.DynamoAttributeValue.fromString("workspace"),
81
      },
82
      updateExpression: "set #status = :error",
83
      expressionAttributeNames: {
84
        "#status": "status",
85
      },
86
      expressionAttributeValues: {
87
        ":error": tasks.DynamoAttributeValue.fromString("error"),
88
      },
89
    }).next(
90
      new sfn.Fail(this, "Fail", {
91
        cause: "Workspace creation failed",
92
      })
93
    );
94

95
    const setCreating = new tasks.DynamoUpdateItem(this, "SetCreating", {
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("creating"),
109
      },
110
      resultPath: sfn.JsonPath.DISCARD,
111
    });
112

113
    const setReady = new tasks.DynamoUpdateItem(this, "SetReady", {
114
      table: props.ragDynamoDBTables.workspacesTable,
115
      key: {
116
        workspace_id: tasks.DynamoAttributeValue.fromString(
117
          sfn.JsonPath.stringAt("$.workspace_id")
118
        ),
119
        object_type: tasks.DynamoAttributeValue.fromString("workspace"),
120
      },
121
      updateExpression: "set #status=:statusValue",
122
      expressionAttributeNames: {
123
        "#status": "status",
124
      },
125
      expressionAttributeValues: {
126
        ":statusValue": tasks.DynamoAttributeValue.fromString("ready"),
127
      },
128
      resultPath: sfn.JsonPath.DISCARD,
129
    });
130

131
    const createTask = new tasks.LambdaInvoke(this, "Create", {
132
      lambdaFunction: createFunction,
133
      resultPath: "$.createResult",
134
    }).addCatch(handleError, {
135
      errors: ["States.ALL"],
136
      resultPath: "$.createResult",
137
    });
138

139
    const workflow = setCreating
140
      .next(createTask)
141
      .next(setReady)
142
      .next(new sfn.Succeed(this, "Success"));
143

144
    const logGroup = new logs.LogGroup(
145
      this,
146
      "CreateOpenSearchWorkspaceSMLogGroup",
147
      {
148
        removalPolicy: RemovalPolicy.DESTROY,
149
      }
150
    );
151

152
    const stateMachine = new sfn.StateMachine(
153
      this,
154
      "CreateOpenSearchWorkspace",
155
      {
156
        definitionBody: sfn.DefinitionBody.fromChainable(workflow),
157
        timeout: cdk.Duration.minutes(5),
158
        comment: "Create OpenSearch Workspace Workflow",
159
        tracingEnabled: true,
160
        logs: {
161
          destination: logGroup,
162
          level: sfn.LogLevel.ALL,
163
        },
164
      }
165
    );
166

167
    this.stateMachine = stateMachine;
168
    this.createWorkspaceRole = createFunction.role;
169
  }
170
}
171