aws-genai-llm-chatbot

1
import * as ec2 from "aws-cdk-lib/aws-ec2";
2
import * as oss from "aws-cdk-lib/aws-opensearchserverless";
3
import * as sfn from "aws-cdk-lib/aws-stepfunctions";
4
import * as iam from "aws-cdk-lib/aws-iam";
5
import { Construct } from "constructs";
6
import { Shared } from "../../shared";
7
import { SystemConfig } from "../../shared/types";
8
import { Utils } from "../../shared/utils";
9
import { RagDynamoDBTables } from "../rag-dynamodb-tables";
10
import { CreateOpenSearchWorkspace } from "./create-opensearch-workspace";
11

12
export interface OpenSearchVectorProps {
13
  readonly config: SystemConfig;
14
  readonly shared: Shared;
15
  readonly ragDynamoDBTables: RagDynamoDBTables;
16
}
17

18
export class OpenSearchVector extends Construct {
19
  public readonly openSearchCollectionName: string;
20
  public readonly openSearchCollectionEndpoint: string;
21
  public readonly openSearchCollection: oss.CfnCollection;
22
  public readonly createOpenSearchWorkspaceWorkflow: sfn.StateMachine;
23
  public addToAccessPolicy: (
24
    name: string,
25
    principal: (string | undefined)[],
26
    permission: string[]
27
  ) => void;
28

29
  constructor(scope: Construct, id: string, props: OpenSearchVectorProps) {
30
    super(scope, id);
31

32
    const collectionName = Utils.getName(
33
      props.config,
34
      "genaichatbot-workspaces"
35
    );
36

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

41
    sg.addIngressRule(
42
      ec2.Peer.ipv4(props.shared.vpc.vpcCidrBlock),
43
      ec2.Port.tcp(443)
44
    );
45

46
    const cfnVpcEndpoint = new oss.CfnVpcEndpoint(this, "VpcEndpoint", {
47
      name: Utils.getName(props.config, "genaichatbot-vpce"),
48
      // Make sure the subnets are not in the same availability zone.
49
      subnetIds: props.shared.vpc.selectSubnets({
50
        onePerAz: true,
51
        subnetType: ec2.SubnetType.PRIVATE_ISOLATED,
52
      }).subnetIds,
53
      vpcId: props.shared.vpc.vpcId,
54
      securityGroupIds: [sg.securityGroupId],
55
    });
56

57
    const cfnNetworkSecurityPolicy = new oss.CfnSecurityPolicy(
58
      this,
59
      "NetworkSecurityPolicy",
60
      {
61
        name: Utils.getName(props.config, "genaichatbot-network-policy"),
62
        type: "network",
63
        policy: JSON.stringify([
64
          {
65
            Rules: [
66
              {
67
                ResourceType: "collection",
68
                Resource: [`collection/${collectionName}`],
69
              },
70
            ],
71
            AllowFromPublic: false,
72
            SourceVPCEs: [cfnVpcEndpoint.attrId],
73
          },
74
        ]).replace(/(\r\n|\n|\r)/gm, ""),
75
      }
76
    );
77

78
    cfnNetworkSecurityPolicy.node.addDependency(cfnVpcEndpoint);
79

80
    const cfnEncryptionSecurityPolicy = new oss.CfnSecurityPolicy(
81
      this,
82
      "EncryptionSecurityPolicy",
83
      {
84
        name: Utils.getName(props.config, "genaichatbot-encryption-policy", 32),
85
        type: "encryption",
86
        policy: JSON.stringify({
87
          Rules: [
88
            {
89
              ResourceType: "collection",
90
              Resource: [`collection/${collectionName}`],
91
            },
92
          ],
93
          AWSOwnedKey: true,
94
        }).replace(/(\r\n|\n|\r)/gm, ""),
95
      }
96
    );
97

98
    cfnEncryptionSecurityPolicy.node.addDependency(cfnNetworkSecurityPolicy);
99

100
    const cfnCollection = new oss.CfnCollection(this, "OpenSearchCollection", {
101
      name: collectionName,
102
      type: "VECTORSEARCH",
103
    });
104

105
    const createWorkflow = new CreateOpenSearchWorkspace(
106
      this,
107
      "CreateOpenSearchWorkspace",
108
      {
109
        config: props.config,
110
        shared: props.shared,
111
        ragDynamoDBTables: props.ragDynamoDBTables,
112
        openSearchCollectionName: collectionName,
113
        openSearchCollection: cfnCollection,
114
        collectionEndpoint: cfnCollection.attrCollectionEndpoint,
115
      }
116
    );
117

118
    cfnCollection.node.addDependency(cfnNetworkSecurityPolicy);
119
    cfnCollection.node.addDependency(cfnEncryptionSecurityPolicy);
120

121
    this.addToAccessPolicyIntl(
122
      props.config,
123
      collectionName,
124
      "create-workflow",
125
      [createWorkflow.createWorkspaceRole?.roleArn],
126
      [
127
        "aoss:CreateIndex",
128
        "aoss:DeleteIndex",
129
        "aoss:UpdateIndex",
130
        "aoss:DescribeIndex",
131
      ]
132
    );
133

134
    this.addToAccessPolicy = (
135
      name: string,
136
      principal: (string | undefined)[],
137
      permission: string[]
138
    ) => {
139
      this.addToAccessPolicyIntl(
140
        props.config,
141
        collectionName,
142
        name,
143
        principal,
144
        permission
145
      );
146
    };
147

148
    this.createOpenSearchWorkspaceWorkflow = createWorkflow.stateMachine;
149
    this.openSearchCollectionEndpoint = cfnCollection.attrCollectionEndpoint;
150
    this.openSearchCollection = cfnCollection;
151
  }
152

153
  private addToAccessPolicyIntl(
154
    config: SystemConfig,
155
    collectionName: string,
156
    name: string,
157
    principal: (string | undefined)[],
158
    permission: string[]
159
  ) {
160
    new oss.CfnAccessPolicy(this, `AccessPolicy-${name}`, {
161
      name: Utils.getName(config, `access-policy-${name}`, 32),
162
      type: "data",
163
      policy: JSON.stringify([
164
        {
165
          Rules: [
166
            {
167
              ResourceType: "index",
168
              Resource: [`index/${collectionName}/*`],
169
              Permission: permission,
170
            },
171
          ],
172
          Principal: principal,
173
        },
174
      ]).replace(/(\r\n|\n|\r)/gm, ""),
175
    });
176
  }
177
}
178