├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── agents ├── app_services_agent_openapi_schema.json ├── multi-tenant-agents.ipynb ├── pool_agent_lambda_function.py └── silo_agent_lambda_function.py ├── cdk ├── .DS_Store ├── .gitignore ├── .npmignore ├── README.md ├── bin │ └── saas-genai-workshop.ts ├── cdk.json ├── jest.config.js ├── lib │ ├── .DS_Store │ ├── control-plane-stack.ts │ ├── core-utils-template-stack.ts │ ├── cp-custom-cognito-auth │ │ ├── auth-custom-resource │ │ │ ├── index.py │ │ │ └── requirements.txt │ │ └── custom-cognito-auth.ts │ ├── destory-policy-setter.ts │ ├── interfaces │ │ └── identity-details.ts │ └── tenant-template │ │ ├── .DS_Store │ │ ├── api-gateway.ts │ │ ├── bedrock-custom │ │ ├── bedrock_logs.py │ │ └── requirements.txt │ │ ├── bedrock.ts │ │ ├── bootstrap-template-stack.ts │ │ ├── cost-per-tenant.ts │ │ ├── cur-athena.ts │ │ ├── cur-report-upload.ts │ │ ├── identity-provider.ts │ │ ├── services.ts │ │ ├── services │ │ ├── .DS_Store │ │ ├── aggregate-metrics │ │ │ ├── invoke_model_tenant_cost.py │ │ │ └── tenant_cost_calculator.py │ │ ├── authorizerService │ │ │ ├── assume_role_layer.py │ │ │ ├── authorizer_layer.py │ │ │ └── tenant_authorizer.py │ │ ├── layers │ │ │ ├── metrics_manager.py │ │ │ └── requirements.txt │ │ ├── ragService │ │ │ └── rag_service.py │ │ ├── s3Uploader │ │ │ └── s3uploader.py │ │ ├── tenant-token-usage │ │ │ └── tenant_token_usage_calculator.py │ │ └── triggerDataIngestionService │ │ │ ├── assume_role_layer.py │ │ │ └── trigger_data_ingestion.py │ │ ├── tenant-provisioning │ │ ├── requirements.txt │ │ └── tenant_provisioning_service.py │ │ ├── tenant-token-usage.ts │ │ ├── usage-plans.ts │ │ └── user-management │ │ └── user_management_service.py ├── package-lock.json ├── package.json ├── test │ └── cdk.test.ts └── tsconfig.json ├── data ├── cur_report.zip ├── tenant1-meeting-notes.txt └── tenant2-meeting-notes.txt ├── scripts ├── .DS_Store ├── cleanup.sh ├── deploy.sh ├── get_opensearch_indices.py ├── install.sh ├── labs │ ├── lab1_add_tenant_provision_service.py │ ├── lab1_provision_tenant_resources.py │ ├── lab2_add_principal_tag.py │ ├── lab2_hardcode_knowledgebase_id.py │ ├── lab2_remove_hardcode_knowledgebase_id.py │ ├── lab3_enable_tenant_token_usage.py │ ├── lab4_calculate_input_outputs_tokens_attribution.py │ ├── lab4_calculate_kb_input_tokens_attribution.py │ ├── lab4_calculate_tenant_cost_kb.py │ ├── lab4_calculate_tenant_cost_service.py │ ├── lab4_get_total_service_cost.py │ └── replace_code.py ├── provision-tenant.sh └── sbt-aws.sh └── vscode-server.yaml /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | ## Code of Conduct 2 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct). 3 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact 4 | opensource-codeofconduct@amazon.com with any additional questions or comments. 5 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing Guidelines 2 | 3 | Thank you for your interest in contributing to our project. Whether it's a bug report, new feature, correction, or additional 4 | documentation, we greatly value feedback and contributions from our community. 5 | 6 | Please read through this document before submitting any issues or pull requests to ensure we have all the necessary 7 | information to effectively respond to your bug report or contribution. 8 | 9 | 10 | ## Reporting Bugs/Feature Requests 11 | 12 | We welcome you to use the GitHub issue tracker to report bugs or suggest features. 13 | 14 | When filing an issue, please check existing open, or recently closed, issues to make sure somebody else hasn't already 15 | reported the issue. Please try to include as much information as you can. Details like these are incredibly useful: 16 | 17 | * A reproducible test case or series of steps 18 | * The version of our code being used 19 | * Any modifications you've made relevant to the bug 20 | * Anything unusual about your environment or deployment 21 | 22 | 23 | ## Contributing via Pull Requests 24 | Contributions via pull requests are much appreciated. Before sending us a pull request, please ensure that: 25 | 26 | 1. You are working against the latest source on the *main* branch. 27 | 2. You check existing open, and recently merged, pull requests to make sure someone else hasn't addressed the problem already. 28 | 3. You open an issue to discuss any significant work - we would hate for your time to be wasted. 29 | 30 | To send us a pull request, please: 31 | 32 | 1. Fork the repository. 33 | 2. Modify the source; please focus on the specific change you are contributing. If you also reformat all the code, it will be hard for us to focus on your change. 34 | 3. Ensure local tests pass. 35 | 4. Commit to your fork using clear commit messages. 36 | 5. Send us a pull request, answering any default questions in the pull request interface. 37 | 6. Pay attention to any automated CI failures reported in the pull request, and stay involved in the conversation. 38 | 39 | GitHub provides additional document on [forking a repository](https://help.github.com/articles/fork-a-repo/) and 40 | [creating a pull request](https://help.github.com/articles/creating-a-pull-request/). 41 | 42 | 43 | ## Finding contributions to work on 44 | Looking at the existing issues is a great way to find something to contribute on. As our projects, by default, use the default GitHub issue labels (enhancement/bug/duplicate/help wanted/invalid/question/wontfix), looking at any 'help wanted' issues is a great place to start. 45 | 46 | 47 | ## Code of Conduct 48 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct). 49 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact 50 | opensource-codeofconduct@amazon.com with any additional questions or comments. 51 | 52 | 53 | ## Security issue notifications 54 | If you discover a potential security issue in this project we ask that you notify AWS/Amazon Security via our [vulnerability reporting page](http://aws.amazon.com/security/vulnerability-reporting/). Please do **not** create a public github issue. 55 | 56 | 57 | ## Licensing 58 | 59 | See the [LICENSE](LICENSE) file for our project's licensing. We will ask you to confirm the licensing of your contribution. 60 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT No Attribution 2 | 3 | Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | this software and associated documentation files (the "Software"), to deal in 7 | the Software without restriction, including without limitation the rights to 8 | use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | the Software, and to permit persons to whom the Software is furnished to do so. 10 | 11 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 12 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 13 | FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 14 | COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 15 | IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 16 | CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 17 | 18 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## SaaS GenAI RAG Workshop 2 | 3 | Follow this link for detailed instructions to run this workshop: https://catalog.us-east-1.prod.workshops.aws/workshops/224f4cc0-cefb-4e29-95fa-365ad5a7ef28 4 | 5 | ## Security 6 | 7 | See [CONTRIBUTING](CONTRIBUTING.md#security-issue-notifications) for more information. 8 | 9 | ## License 10 | 11 | This library is licensed under the MIT-0 License. See the LICENSE file. 12 | -------------------------------------------------------------------------------- /agents/app_services_agent_openapi_schema.json: -------------------------------------------------------------------------------- 1 | { 2 | "openapi": "3.0.0", 3 | "info": { 4 | "title": "Application Services Automation API", 5 | "version": "1.0.0", 6 | "description": "APIs for creating products and orders." 7 | }, 8 | "components": { 9 | "schemas": { 10 | "Product": { 11 | "type": "object", 12 | "properties": { 13 | "productId": { 14 | "type": "string", 15 | "description": "Unique ID of the product." 16 | }, 17 | "productName": { 18 | "type": "string", 19 | "description": "product name." 20 | }, 21 | "productPrice": { 22 | "type": "string", 23 | "description": "product price." 24 | } 25 | } 26 | }, 27 | "Order": { 28 | "type": "object", 29 | "properties": { 30 | "orderId": { 31 | "type": "string", 32 | "description": "Unique ID of the order." 33 | }, 34 | "quantity": { 35 | "type": "string", 36 | "description": "Number of products to be ordered." 37 | }, 38 | "product": { 39 | "$ref": "#/components/schemas/Product" 40 | } 41 | } 42 | } 43 | } 44 | }, 45 | "paths": { 46 | "/products/create-product": { 47 | "post": { 48 | "summary": "API to create a product", 49 | "description": "create a product. The API takes product name and product price.", 50 | "operationId": "createProduct", 51 | "requestBody": { 52 | "required": true, 53 | "content": { 54 | "application/json": { 55 | "schema": { 56 | "type": "object", 57 | "properties": { 58 | "productName": { 59 | "type": "string", 60 | "description": "product name." 61 | }, 62 | "productPrice": { 63 | "type": "string", 64 | "description": "product price." 65 | } 66 | }, 67 | "required": ["productName", "productPrice"] 68 | } 69 | } 70 | } 71 | }, 72 | "responses": { 73 | "200": { 74 | "description": "Product created successfully", 75 | "content": { 76 | "application/json": { 77 | "schema": { 78 | "$ref": "#/components/schemas/Product" 79 | } 80 | } 81 | } 82 | }, 83 | "400": { 84 | "description": "Bad request. One or more required fields are missing or invalid." 85 | } 86 | } 87 | } 88 | }, 89 | "/products": { 90 | "get": { 91 | "summary": "API to gets all products", 92 | "description": "Gets all product details. The API gets all product details.", 93 | "operationId": "getProducts", 94 | "responses": { 95 | "200": { 96 | "description": "Products retrieved successfully", 97 | "content": { 98 | "application/json": { 99 | "schema": { 100 | "type": "array", 101 | "items": { 102 | "$ref": "#/components/schemas/Product" 103 | } 104 | } 105 | } 106 | } 107 | } 108 | } 109 | } 110 | }, 111 | "/products/{productId}": { 112 | "get": { 113 | "summary": "API to get a product", 114 | "description": "Gets single product details. The API takes productId and gets the product details.", 115 | "operationId": "getProduct", 116 | "parameters": [ 117 | { 118 | "name": "productId", 119 | "in": "path", 120 | "description": "Unique ID of the product", 121 | "required": true, 122 | "schema": { 123 | "type": "string" 124 | } 125 | } 126 | ], 127 | "responses": { 128 | "200": { 129 | "description": "Products retrieved successfully", 130 | "content": { 131 | "application/json": { 132 | "schema": { 133 | "$ref": "#/components/schemas/Product" 134 | } 135 | } 136 | } 137 | }, 138 | "400": { 139 | "description": "Bad request. One or more required fields are missing or invalid." 140 | } 141 | } 142 | } 143 | }, 144 | "/orders/create-order": { 145 | "post": { 146 | "summary": "API to create a order", 147 | "description": "create a order. The API takes a product, quantity and orders the product.", 148 | "operationId": "createOrder", 149 | "requestBody": { 150 | "required": true, 151 | "content": { 152 | "application/json": { 153 | "schema": { 154 | "type": "object", 155 | "properties": { 156 | "product": { 157 | "$ref": "#/components/schemas/Product" 158 | }, 159 | "quantity": { 160 | "type": "string", 161 | "description": "Number of products to be ordered." 162 | } 163 | }, 164 | "required": ["product", "quantity"] 165 | } 166 | } 167 | } 168 | }, 169 | "responses": { 170 | "200": { 171 | "description": "Order created successfully", 172 | "content": { 173 | "application/json": { 174 | "schema": { 175 | "$ref": "#/components/schemas/Order" 176 | } 177 | } 178 | } 179 | }, 180 | "400": { 181 | "description": "Bad request. One or more required fields are missing or invalid." 182 | } 183 | } 184 | } 185 | }, 186 | "/orders": { 187 | "get": { 188 | "summary": "API to gets all orders", 189 | "description": "Gets all orders. The API get all order details.", 190 | "operationId": "getOrders", 191 | "responses": { 192 | "200": { 193 | "description": "Orders retrieved successfully", 194 | "content": { 195 | "application/json": { 196 | "schema": { 197 | "type": "array", 198 | "items": { 199 | "$ref": "#/components/schemas/Order" 200 | } 201 | } 202 | } 203 | } 204 | } 205 | } 206 | } 207 | } 208 | } 209 | } 210 | -------------------------------------------------------------------------------- /agents/pool_agent_lambda_function.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 3 | # SPDX-License-Identifier: MIT-0 4 | import json 5 | import boto3 6 | import uuid 7 | from boto3.dynamodb.conditions import Key 8 | import os 9 | 10 | def get_named_parameter(event, name): 11 | return next(item for item in event['parameters'] if item['name'] == name)['value'] 12 | 13 | 14 | def get_named_property(event, name): 15 | return next( 16 | item for item in 17 | event['requestBody']['content']['application/json']['properties'] 18 | if item['name'] == name)['value'] 19 | 20 | 21 | 22 | 23 | def create_product(event, productName, productPrice, tenantId): 24 | print('create product invoked') 25 | productId=uuid.uuid4().hex 26 | table = __get_dynamodb_table(event, 'pool-product-table') 27 | response = table.put_item( 28 | Item= 29 | { 30 | 'tenantId': tenantId, 31 | 'productId': productId, 32 | 'productName': productName, 33 | 'productPrice': productPrice 34 | } 35 | ) 36 | return { 37 | "response": { 38 | 'productId': productId, 39 | 'productName': productName, 40 | 'productPrice': productPrice 41 | } 42 | } 43 | 44 | 45 | def get_products(event, tenantId): 46 | print('get products invoked') 47 | get_all_response =[] 48 | table = __get_dynamodb_table(event, 'pool-product-table') 49 | 50 | __get_tenant_data(tenantId,get_all_response, table) 51 | return { 52 | "response": get_all_response 53 | } 54 | 55 | 56 | 57 | def get_product(event, productId, tenantId): 58 | print('get product invoked') 59 | table = __get_dynamodb_table(event, 'pool-product-table') 60 | response =table.get_item(Key={'tenantId': tenantId, 'productId': productId}) 61 | item = response['Item'] 62 | return { 63 | "response": { 64 | 'productId': item['productId'], 65 | 'productName': item['productName'], 66 | 'productPrice': item['productPrice'] 67 | } 68 | } 69 | 70 | def create_order(event, product, tenantId, quantity): 71 | print('create order invoked') 72 | orderId=uuid.uuid4().hex 73 | table = __get_dynamodb_table(event, 'pool-order-table') 74 | response = table.put_item( 75 | Item= 76 | { 77 | 'tenantId': tenantId, 78 | 'orderId': orderId, 79 | 'quantity': quantity, 80 | 'product': product 81 | } 82 | ) 83 | return { 84 | "response": { 85 | 'orderId': orderId, 86 | 'quantity': quantity, 87 | 'product': product 88 | } 89 | } 90 | 91 | def get_orders(event, tenantId): 92 | print('get orders invoked') 93 | get_all_response =[] 94 | table = __get_dynamodb_table(event, 'pool-order-table') 95 | 96 | __get_tenant_data(tenantId,get_all_response, table) 97 | return { 98 | "response": get_all_response 99 | } 100 | 101 | 102 | def __get_tenant_data(tenant_id, get_all_response, table): 103 | response = table.query(KeyConditionExpression=Key('tenantId').eq(tenant_id)) 104 | if (len(response['Items']) > 0): 105 | for item in response['Items']: 106 | get_all_response.append(item) 107 | 108 | 109 | 110 | def lambda_handler(event, context): 111 | action = event['actionGroup'] 112 | api_path = event['apiPath'] 113 | tenantId = event['sessionAttributes']['tenantId'] 114 | 115 | if api_path == '/products/create-product': 116 | productName = get_named_property(event, "productName") 117 | productPrice = get_named_property(event, "productPrice") 118 | body = create_product(event, productName, productPrice, tenantId) 119 | elif api_path == '/products': 120 | body = get_products(event, tenantId) 121 | elif api_path == '/products/{productId}': 122 | productId = get_named_parameter(event, "productId") 123 | body = get_product(event, productId, tenantId) 124 | elif api_path == '/orders/create-order': 125 | product = get_named_property(event, "product") 126 | quantity = get_named_property(event, "quantity") 127 | body = create_order(event, product, tenantId,quantity) 128 | elif api_path == '/orders': 129 | body = get_orders(event, tenantId) 130 | else: 131 | body = {"{}::{} is not a valid api, try another one.".format(action, api_path)} 132 | 133 | response_body = { 134 | 'application/json': { 135 | 'body': str(body) 136 | } 137 | } 138 | 139 | action_response = { 140 | 'actionGroup': event['actionGroup'], 141 | 'apiPath': event['apiPath'], 142 | 'httpMethod': event['httpMethod'], 143 | 'httpStatusCode': 200, 144 | 'responseBody': response_body 145 | } 146 | 147 | response = {'response': action_response} 148 | return response 149 | 150 | def __get_dynamodb_table(event, table_name): 151 | dynamodb = boto3.resource('dynamodb', 152 | aws_access_key_id=event['sessionAttributes']['accessKeyId'], 153 | aws_secret_access_key=event['sessionAttributes']['secretAccessKey'], 154 | aws_session_token=event['sessionAttributes']['sessionToken'] 155 | ) 156 | 157 | return dynamodb.Table(table_name) 158 | 159 | 160 | -------------------------------------------------------------------------------- /agents/silo_agent_lambda_function.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 3 | # SPDX-License-Identifier: MIT-0 4 | import json 5 | import boto3 6 | import uuid 7 | from boto3.dynamodb.conditions import Key 8 | import os 9 | 10 | def get_named_parameter(event, name): 11 | return next(item for item in event['parameters'] if item['name'] == name)['value'] 12 | 13 | 14 | def get_named_property(event, name): 15 | return next( 16 | item for item in 17 | event['requestBody']['content']['application/json']['properties'] 18 | if item['name'] == name)['value'] 19 | 20 | 21 | 22 | 23 | def create_product(event, productName, productPrice, tenantId): 24 | print('create product invoked') 25 | productId=uuid.uuid4().hex 26 | table = __get_dynamodb_table(event, 'silo-product-table') 27 | response = table.put_item( 28 | Item= 29 | { 30 | 'tenantId': tenantId, 31 | 'productId': productId, 32 | 'productName': productName, 33 | 'productPrice': productPrice 34 | } 35 | ) 36 | return { 37 | "response": { 38 | 'productId': productId, 39 | 'productName': productName, 40 | 'productPrice': productPrice 41 | } 42 | } 43 | 44 | 45 | def get_products(event, tenantId): 46 | print('get products invoked') 47 | get_all_response =[] 48 | table = __get_dynamodb_table(event, 'silo-product-table') 49 | 50 | __get_tenant_data(tenantId,get_all_response, table) 51 | return { 52 | "response": get_all_response 53 | } 54 | 55 | 56 | 57 | def get_product(event, productId, tenantId): 58 | print('get product invoked') 59 | table = __get_dynamodb_table(event, 'silo-product-table') 60 | response =table.get_item(Key={'tenantId': tenantId, 'productId': productId}) 61 | item = response['Item'] 62 | return { 63 | "response": { 64 | 'productId': item['productId'], 65 | 'productName': item['productName'], 66 | 'productPrice': item['productPrice'] 67 | } 68 | } 69 | 70 | def create_order(event, product, tenantId, quantity): 71 | print('create order invoked') 72 | orderId=uuid.uuid4().hex 73 | table = __get_dynamodb_table(event, 'silo-order-table') 74 | response = table.put_item( 75 | Item= 76 | { 77 | 'tenantId': tenantId, 78 | 'orderId': orderId, 79 | 'quantity': quantity, 80 | 'product': product 81 | } 82 | ) 83 | return { 84 | "response": { 85 | 'orderId': orderId, 86 | 'quantity': quantity, 87 | 'product': product 88 | } 89 | } 90 | 91 | def get_orders(event, tenantId): 92 | print('get orders invoked') 93 | get_all_response =[] 94 | table = __get_dynamodb_table(event, 'silo-order-table') 95 | 96 | __get_tenant_data(tenantId,get_all_response, table) 97 | return { 98 | "response": get_all_response 99 | } 100 | 101 | 102 | def __get_tenant_data(tenant_id, get_all_response, table): 103 | response = table.query(KeyConditionExpression=Key('tenantId').eq(tenant_id)) 104 | if (len(response['Items']) > 0): 105 | for item in response['Items']: 106 | get_all_response.append(item) 107 | 108 | 109 | 110 | def lambda_handler(event, context): 111 | action = event['actionGroup'] 112 | api_path = event['apiPath'] 113 | tenantId = event['sessionAttributes']['tenantId'] 114 | 115 | if api_path == '/products/create-product': 116 | productName = get_named_property(event, "productName") 117 | productPrice = get_named_property(event, "productPrice") 118 | body = create_product(event, productName, productPrice, tenantId) 119 | elif api_path == '/products': 120 | body = get_products(event, tenantId) 121 | elif api_path == '/products/{productId}': 122 | productId = get_named_parameter(event, "productId") 123 | body = get_product(event, productId, tenantId) 124 | elif api_path == '/orders/create-order': 125 | product = get_named_property(event, "product") 126 | quantity = get_named_property(event, "quantity") 127 | body = create_order(event, product, tenantId,quantity) 128 | elif api_path == '/orders': 129 | body = get_orders(event, tenantId) 130 | else: 131 | body = {"{}::{} is not a valid api, try another one.".format(action, api_path)} 132 | 133 | response_body = { 134 | 'application/json': { 135 | 'body': str(body) 136 | } 137 | } 138 | 139 | action_response = { 140 | 'actionGroup': event['actionGroup'], 141 | 'apiPath': event['apiPath'], 142 | 'httpMethod': event['httpMethod'], 143 | 'httpStatusCode': 200, 144 | 'responseBody': response_body 145 | } 146 | 147 | response = {'response': action_response} 148 | return response 149 | 150 | def __get_dynamodb_table(event, table_name): 151 | dynamodb = boto3.resource('dynamodb') 152 | 153 | return dynamodb.Table(table_name) 154 | 155 | 156 | -------------------------------------------------------------------------------- /cdk/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/aws-saas-genai-rag-workshop/e81b3a60cdb6adaf4d1f9c1fb11e279b279d473d/cdk/.DS_Store -------------------------------------------------------------------------------- /cdk/.gitignore: -------------------------------------------------------------------------------- 1 | *.js 2 | !jest.config.js 3 | *.d.ts 4 | node_modules 5 | 6 | # CDK asset staging directory 7 | .cdk.staging 8 | cdk.out 9 | -------------------------------------------------------------------------------- /cdk/.npmignore: -------------------------------------------------------------------------------- 1 | *.ts 2 | !*.d.ts 3 | 4 | # CDK asset staging directory 5 | .cdk.staging 6 | cdk.out 7 | -------------------------------------------------------------------------------- /cdk/README.md: -------------------------------------------------------------------------------- 1 | # Welcome to your CDK TypeScript project 2 | 3 | This is a blank project for CDK development with TypeScript. 4 | 5 | The `cdk.json` file tells the CDK Toolkit how to execute your app. 6 | 7 | 8 | ## Useful commands 9 | 10 | * `npm run build` compile typescript to js 11 | * `npm run watch` watch for changes and compile 12 | * `npm run test` perform the jest unit tests 13 | * `npx cdk deploy` deploy this stack to your default AWS account/region 14 | * `npx cdk diff` compare deployed stack with current state 15 | * `npx cdk synth` emits the synthesized CloudFormation template 16 | -------------------------------------------------------------------------------- /cdk/bin/saas-genai-workshop.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import "source-map-support/register"; 5 | import * as cdk from "aws-cdk-lib"; 6 | import { BootstrapTemplateStack } from "../lib/tenant-template/bootstrap-template-stack"; 7 | import { ControlPlaneStack } from "../lib/control-plane-stack"; 8 | import { CoreUitlsTemplateStack } from "../lib/core-utils-template-stack"; 9 | 10 | const app = new cdk.App(); 11 | 12 | const controlPlaneStack = new ControlPlaneStack(app, "ControlPlaneStack", { 13 | systemAdminRoleName: process.env.CDK_PARAM_SYSTEM_ADMIN_ROLE_NAME, 14 | systemAdminEmail: process.env.CDK_PARAM_SYSTEM_ADMIN_EMAIL, 15 | }); 16 | 17 | const CoreUtilsTemplateStack = new CoreUitlsTemplateStack( 18 | app, 19 | "saas-genai-workshop-core-utils-stack", 20 | { 21 | controlPlane: controlPlaneStack.controlPlane, 22 | } 23 | ); 24 | 25 | const bootstrapTemplateStack = new BootstrapTemplateStack( 26 | app, 27 | "saas-genai-workshop-bootstrap-template", 28 | { 29 | coreUtilsStack: CoreUtilsTemplateStack, 30 | controlPlaneApiGwUrl: 31 | controlPlaneStack.controlPlane.controlPlaneAPIGatewayUrl, 32 | } 33 | ); 34 | -------------------------------------------------------------------------------- /cdk/cdk.json: -------------------------------------------------------------------------------- 1 | { 2 | "app": "npx ts-node --prefer-ts-exts bin/saas-genai-workshop.ts", 3 | "watch": { 4 | "include": ["**"], 5 | "exclude": [ 6 | "README.md", 7 | "cdk*.json", 8 | "**/*.d.ts", 9 | "**/*.js", 10 | "tsconfig.json", 11 | "package*.json", 12 | "yarn.lock", 13 | "node_modules", 14 | "test" 15 | ] 16 | }, 17 | "context": { 18 | "@aws-cdk/aws-lambda:recognizeLayerVersion": true, 19 | "@aws-cdk/core:checkSecretUsage": true, 20 | "@aws-cdk/core:target-partitions": ["aws", "aws-cn"], 21 | "@aws-cdk-containers/ecs-service-extensions:enableDefaultLogDriver": true, 22 | "@aws-cdk/aws-ec2:uniqueImdsv2TemplateName": true, 23 | "@aws-cdk/aws-ecs:arnFormatIncludesClusterName": true, 24 | "@aws-cdk/aws-iam:minimizePolicies": true, 25 | "@aws-cdk/core:validateSnapshotRemovalPolicy": true, 26 | "@aws-cdk/aws-codepipeline:crossAccountKeyAliasStackSafeResourceName": true, 27 | "@aws-cdk/aws-s3:createDefaultLoggingPolicy": true, 28 | "@aws-cdk/aws-sns-subscriptions:restrictSqsDescryption": true, 29 | "@aws-cdk/aws-apigateway:disableCloudWatchRole": true, 30 | "@aws-cdk/core:enablePartitionLiterals": true, 31 | "@aws-cdk/aws-events:eventsTargetQueueSameAccount": true, 32 | "@aws-cdk/aws-iam:standardizedServicePrincipals": true, 33 | "@aws-cdk/aws-ecs:disableExplicitDeploymentControllerForCircuitBreaker": true, 34 | "@aws-cdk/aws-iam:importedRoleStackSafeDefaultPolicyName": true, 35 | "@aws-cdk/aws-s3:serverAccessLogsUseBucketPolicy": true, 36 | "@aws-cdk/aws-route53-patters:useCertificate": true, 37 | "@aws-cdk/customresources:installLatestAwsSdkDefault": false, 38 | "@aws-cdk/aws-rds:databaseProxyUniqueResourceName": true, 39 | "@aws-cdk/aws-codedeploy:removeAlarmsFromDeploymentGroup": true, 40 | "@aws-cdk/aws-apigateway:authorizerChangeDeploymentLogicalId": true, 41 | "@aws-cdk/aws-ec2:launchTemplateDefaultUserData": true, 42 | "@aws-cdk/aws-secretsmanager:useAttachedSecretResourcePolicyForSecretTargetAttachments": true, 43 | "@aws-cdk/aws-redshift:columnId": true, 44 | "@aws-cdk/aws-stepfunctions-tasks:enableEmrServicePolicyV2": true, 45 | "@aws-cdk/aws-ec2:restrictDefaultSecurityGroup": true, 46 | "@aws-cdk/aws-apigateway:requestValidatorUniqueId": true, 47 | "@aws-cdk/aws-kms:aliasNameRef": true, 48 | "@aws-cdk/aws-autoscaling:generateLaunchTemplateInsteadOfLaunchConfig": true, 49 | "@aws-cdk/core:includePrefixInUniqueNameGeneration": true, 50 | "@aws-cdk/aws-efs:denyAnonymousAccess": true, 51 | "@aws-cdk/aws-opensearchservice:enableOpensearchMultiAzWithStandby": true, 52 | "@aws-cdk/aws-lambda-nodejs:useLatestRuntimeVersion": true, 53 | "@aws-cdk/aws-efs:mountTargetOrderInsensitiveLogicalId": true, 54 | "@aws-cdk/aws-rds:auroraClusterChangeScopeOfInstanceParameterGroupWithEachParameters": true, 55 | "@aws-cdk/aws-appsync:useArnForSourceApiAssociationIdentifier": true, 56 | "@aws-cdk/aws-rds:preventRenderingDeprecatedCredentials": true, 57 | "@aws-cdk/aws-codepipeline-actions:useNewDefaultBranchForCodeCommitSource": true 58 | } 59 | } 60 | -------------------------------------------------------------------------------- /cdk/jest.config.js: -------------------------------------------------------------------------------- 1 | module.exports = { 2 | testEnvironment: 'node', 3 | roots: ['/test'], 4 | testMatch: ['**/*.test.ts'], 5 | transform: { 6 | '^.+\\.tsx?$': 'ts-jest' 7 | } 8 | }; 9 | -------------------------------------------------------------------------------- /cdk/lib/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/aws-saas-genai-rag-workshop/e81b3a60cdb6adaf4d1f9c1fb11e279b279d473d/cdk/lib/.DS_Store -------------------------------------------------------------------------------- /cdk/lib/control-plane-stack.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import { ControlPlane, CognitoAuth } from "@cdklabs/sbt-aws"; 5 | import { Stack, StackProps } from "aws-cdk-lib"; 6 | import { Construct } from "constructs"; 7 | import { CustomCognitoAuth } from "./cp-custom-cognito-auth/custom-cognito-auth"; 8 | 9 | interface ControlPlaneStackProps extends StackProps { 10 | readonly systemAdminRoleName?: string; 11 | readonly systemAdminEmail?: string; 12 | } 13 | 14 | export class ControlPlaneStack extends Stack { 15 | public readonly regApiGatewayUrl: string; 16 | public readonly controlPlane: ControlPlane; 17 | 18 | constructor(scope: Construct, id: string, props: ControlPlaneStackProps) { 19 | super(scope, id, props); 20 | 21 | const systemAdminEmail = 22 | props.systemAdminEmail || 23 | process.env.CDK_PARAM_SYSTEM_ADMIN_EMAIL || 24 | "admin@example.com"; 25 | 26 | if (!process.env.CDK_PARAM_SYSTEM_ADMIN_ROLE_NAME) { 27 | process.env.CDK_PARAM_SYSTEM_ADMIN_ROLE_NAME = "SystemAdmin"; 28 | } 29 | 30 | const customCognitoAuth = new CustomCognitoAuth(this, "CognitoAuth", { 31 | // Avoid checking scopes for API endpoints. Done only for testing purposes. 32 | setAPIGWScopes: false, 33 | }); 34 | 35 | // TODO: Lab1 - Add SBT Control plane 36 | const controlPlane = new ControlPlane(this, "ControlPlane", { 37 | auth: customCognitoAuth, 38 | systemAdminEmail: systemAdminEmail, 39 | }); 40 | this.controlPlane = controlPlane; 41 | 42 | this.regApiGatewayUrl = controlPlane.controlPlaneAPIGatewayUrl; 43 | this.controlPlane = controlPlane; 44 | } 45 | } 46 | -------------------------------------------------------------------------------- /cdk/lib/core-utils-template-stack.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import { Stack, StackProps } from "aws-cdk-lib"; 5 | import { Construct } from "constructs"; 6 | import { PolicyDocument } from "aws-cdk-lib/aws-iam"; 7 | import { Project } from "aws-cdk-lib/aws-codebuild"; 8 | import { 9 | CoreApplicationPlane, 10 | DetailType, 11 | EventManager, 12 | ControlPlane, 13 | BashJobRunnerProps, 14 | BashJobRunner, 15 | } from "@cdklabs/sbt-aws"; 16 | import * as fs from "fs"; 17 | 18 | interface CoreUtilsTemplateStackProps extends StackProps { 19 | readonly controlPlane: ControlPlane; 20 | } 21 | 22 | export class CoreUitlsTemplateStack extends Stack { 23 | public readonly codeBuildProject: Project; 24 | 25 | constructor( 26 | scope: Construct, 27 | id: string, 28 | props: CoreUtilsTemplateStackProps 29 | ) { 30 | super(scope, id, props); 31 | 32 | const provisioningJobRunnerProps: BashJobRunnerProps = { 33 | permissions: PolicyDocument.fromJson( 34 | JSON.parse(` 35 | { 36 | "Version": "2012-10-17", 37 | "Statement": [ 38 | { 39 | "Effect": "Allow", 40 | "Action": [ 41 | "cloudformation:DescribeStacks", 42 | "cognito-idp:AdminCreateUser", 43 | "cognito-idp:AdminSetUserPassword", 44 | "cognito-idp:AdminUpdateUserAttributes", 45 | "cognito-idp:AdminGetUser", 46 | "cognito-idp:CreateGroup", 47 | "cognito-idp:AdminAddUserToGroup", 48 | "cognito-idp:GetGroup", 49 | "s3:PutObject", 50 | "s3:GetObject", 51 | "s3:ListBucket", 52 | "lambda:AddPermission", 53 | "events:PutRule", 54 | "events:PutTargets", 55 | "iam:CreateRole", 56 | "iam:GetRole", 57 | "iam:PutRolePolicy", 58 | "iam:PassRole", 59 | "bedrock:CreateKnowledgeBase", 60 | "bedrock:CreateDataSource", 61 | "bedrock:CreateKnowledgeBase", 62 | "bedrock:InvokeModel", 63 | "bedrock:ListKnowledgeBases", 64 | "aoss:CreateAccessPolicy", 65 | "aoss:BatchGetCollection", 66 | "aoss:APIAccessAll", 67 | "codecommit:GetRepository", 68 | "codecommit:GitPull", 69 | "apigateway:*" 70 | ], 71 | "Resource": "*" 72 | } 73 | ] 74 | } 75 | `) 76 | ), 77 | script: fs.readFileSync("../scripts/provision-tenant.sh", "utf8"), 78 | environmentJSONVariablesFromIncomingEvent: [ 79 | "tenantId", 80 | "tenantName", 81 | "email", 82 | "tenantStatus", 83 | ], 84 | environmentVariablesToOutgoingEvent: ["tenantStatus", "tenantConfig"], 85 | scriptEnvironmentVariables: {}, 86 | outgoingEvent: DetailType.PROVISION_SUCCESS, 87 | incomingEvent: DetailType.ONBOARDING_REQUEST, 88 | eventManager: props.controlPlane.eventManager, 89 | }; 90 | 91 | const provisioningJobRunner: BashJobRunner = new BashJobRunner( 92 | this, 93 | "provisioningJobRunner", 94 | provisioningJobRunnerProps 95 | ); 96 | 97 | this.codeBuildProject = provisioningJobRunner.codebuildProject; 98 | 99 | // TODO: Lab1 - Add SBT core utils component 100 | new CoreApplicationPlane(this, "CoreApplicationPlane", { 101 | eventManager: props.controlPlane.eventManager, 102 | jobRunnersList: [provisioningJobRunner], 103 | }); 104 | } 105 | } 106 | -------------------------------------------------------------------------------- /cdk/lib/cp-custom-cognito-auth/auth-custom-resource/index.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | import boto3 5 | import botocore 6 | from crhelper import CfnResource 7 | cognito = boto3.client('cognito-idp') 8 | helper = CfnResource() 9 | 10 | 11 | @helper.create 12 | @helper.update 13 | def do_action(event, _): 14 | """ Called as part of bootstrap template. 15 | Inserts/Updates Settings table based upon the resources deployed inside bootstrap template 16 | We use these settings inside tenant template 17 | 18 | Args: 19 | event ([type]): [description] 20 | _ ([type]): [description] 21 | """ 22 | # Setting email address as username 23 | user_name = event['ResourceProperties']['Email'] 24 | email = event['ResourceProperties']['Email'] 25 | user_role = event['ResourceProperties']['Role'] 26 | user_pool_id = event['ResourceProperties']['UserPoolId'] 27 | 28 | try: 29 | create_user_response = cognito.admin_create_user( 30 | Username=user_name, 31 | UserPoolId=user_pool_id, 32 | ForceAliasCreation=True, 33 | MessageAction='SUPPRESS', 34 | UserAttributes=[ 35 | { 36 | 'Name': 'email', 37 | 'Value': email 38 | }, 39 | { 40 | 'Name': 'email_verified', 41 | 'Value': 'true' 42 | }, 43 | { 44 | 'Name': 'custom:userRole', 45 | 'Value': user_role 46 | } 47 | ] 48 | ) 49 | # Set a default password 50 | cognito.admin_set_user_password( 51 | UserPoolId=user_pool_id, 52 | Username=user_name, 53 | Password='SaaS123!', 54 | Permanent=True 55 | ) 56 | 57 | print(f'create_user_response: {create_user_response}') 58 | except cognito.exceptions.UsernameExistsException: 59 | print(f'user: {user_name} already exists!') 60 | return user_name 61 | 62 | 63 | @helper.delete 64 | def do_delete(event, _): 65 | user_name = event["PhysicalResourceId"] 66 | user_pool_id = event['ResourceProperties']['UserPoolId'] 67 | try: 68 | cognito.admin_delete_user( 69 | UserPoolId=user_pool_id, 70 | Username=user_name 71 | ) 72 | except botocore.exceptions.ClientError as e: 73 | print(f'failed to delete: {e}') 74 | 75 | 76 | def handler(event, context): 77 | helper(event, context) 78 | -------------------------------------------------------------------------------- /cdk/lib/cp-custom-cognito-auth/auth-custom-resource/requirements.txt: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | 4 | crhelper -------------------------------------------------------------------------------- /cdk/lib/cp-custom-cognito-auth/custom-cognito-auth.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import * as sbt from "@cdklabs/sbt-aws"; 5 | import { Stack, CustomResource, Duration } from "aws-cdk-lib"; 6 | import * as lambda from "aws-cdk-lib/aws-lambda"; 7 | import * as python from "@aws-cdk/aws-lambda-python-alpha"; 8 | import { Construct } from "constructs"; 9 | import * as path from "path"; 10 | 11 | export class CustomCognitoAuth extends sbt.CognitoAuth { 12 | private readonly customResourceFunction: lambda.IFunction; 13 | 14 | constructor(scope: Construct, id: string, props: sbt.CognitoAuthProps) { 15 | super(scope, id, props); 16 | // https://docs.powertools.aws.dev/lambda/python/2.31.0/#lambda-layer 17 | const lambdaPowertoolsLayer = lambda.LayerVersion.fromLayerVersionArn( 18 | this, 19 | "LambdaPowerToolsCustumAuth", 20 | `arn:aws:lambda:${ 21 | Stack.of(this).region 22 | }:017000801446:layer:AWSLambdaPowertoolsPythonV2:59` 23 | ); 24 | 25 | this.customResourceFunction = new python.PythonFunction( 26 | this, 27 | "customResourceFunction", 28 | { 29 | entry: path.join(__dirname, "auth-custom-resource"), 30 | runtime: lambda.Runtime.PYTHON_3_12, 31 | index: "index.py", 32 | handler: "handler", 33 | timeout: Duration.seconds(60), 34 | layers: [lambdaPowertoolsLayer], 35 | } 36 | ); 37 | this.userPool.grant( 38 | this.customResourceFunction, 39 | "cognito-idp:AdminCreateUser", 40 | "cognito-idp:AdminDeleteUser", 41 | "cognito-idp:AdminSetUserPassword" 42 | ); 43 | } 44 | createAdminUser( 45 | scope: Construct, 46 | id: string, 47 | props: sbt.CreateAdminUserProps 48 | ) { 49 | new CustomResource(scope, `CustomAuthCustomResource-${id}`, { 50 | serviceToken: this.customResourceFunction.functionArn, 51 | properties: { 52 | UserPoolId: this.userPool.userPoolId, 53 | Name: props.name, 54 | Email: props.email, 55 | Role: props.role, 56 | }, 57 | }); 58 | } 59 | } 60 | -------------------------------------------------------------------------------- /cdk/lib/destory-policy-setter.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import { CfnResource, IAspect, RemovalPolicy } from "aws-cdk-lib"; 5 | import { IConstruct } from "constructs"; 6 | 7 | export class DestroyPolicySetter implements IAspect { 8 | public visit(node: IConstruct): void { 9 | if (node instanceof CfnResource) { 10 | node.applyRemovalPolicy(RemovalPolicy.DESTROY); 11 | } 12 | } 13 | } 14 | -------------------------------------------------------------------------------- /cdk/lib/interfaces/identity-details.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | export interface IdentityDetails { 5 | name: string; 6 | details: { 7 | [key: string]: any; 8 | }; 9 | } 10 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/aws-saas-genai-rag-workshop/e81b3a60cdb6adaf4d1f9c1fb11e279b279d473d/cdk/lib/tenant-template/.DS_Store -------------------------------------------------------------------------------- /cdk/lib/tenant-template/api-gateway.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import { Construct } from "constructs"; 5 | import { 6 | RestApi, 7 | Cors, 8 | LogGroupLogDestination, 9 | MethodLoggingLevel, 10 | Period, 11 | ApiKeySourceType, 12 | UsagePlan, 13 | } from "aws-cdk-lib/aws-apigateway"; 14 | import { LogGroup, RetentionDays } from "aws-cdk-lib/aws-logs"; 15 | import { Services } from "./services"; 16 | import { UsagePlans } from "./usage-plans"; 17 | 18 | interface ApiGatewayProps {} 19 | 20 | export class ApiGateway extends Construct { 21 | public readonly restApi: RestApi; 22 | public readonly usagePlaneApiKey: any; 23 | public readonly usagePlanBasicTier: UsagePlan; 24 | 25 | constructor(scope: Construct, id: string, props: ApiGatewayProps) { 26 | super(scope, id); 27 | 28 | const apiLogGroup = new LogGroup(this, "SaaSGenAIWorkshopAPILogGroup", { 29 | retention: RetentionDays.ONE_WEEK, 30 | }); 31 | 32 | const restApi = new RestApi(this, "SaaSGenAIWorkshopRestApi", { 33 | cloudWatchRole: true, 34 | apiKeySourceType: ApiKeySourceType.AUTHORIZER, 35 | defaultCorsPreflightOptions: { 36 | allowOrigins: Cors.ALL_ORIGINS, 37 | }, 38 | deployOptions: { 39 | accessLogDestination: new LogGroupLogDestination(apiLogGroup), 40 | methodOptions: { 41 | "/*/*": { 42 | dataTraceEnabled: true, 43 | loggingLevel: MethodLoggingLevel.ERROR, 44 | }, 45 | }, 46 | }, 47 | }); 48 | 49 | this.restApi = restApi; 50 | 51 | const usagePlan = new UsagePlans(this, "usagePlan", { 52 | apiGateway: restApi, 53 | }); 54 | 55 | this.usagePlanBasicTier = usagePlan.usagePlanBasicTier; 56 | } 57 | } 58 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/bedrock-custom/bedrock_logs.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | import boto3 5 | import os 6 | from crhelper import CfnResource 7 | 8 | log_group = os.environ['LOG_GROUP_NAME'] 9 | log_role= os.environ['BEDROCK_LOG_ROLE'] 10 | 11 | # Create Bedrock client 12 | bedrock_client = boto3.client('bedrock') 13 | helper = CfnResource() 14 | 15 | @helper.create 16 | @helper.update 17 | def do_action(event, _): 18 | # Set up logging configuration 19 | logging_config = { 20 | 'cloudWatchConfig': { 21 | 'logGroupName': log_group, 22 | 'roleArn': log_role 23 | } 24 | 25 | } 26 | 27 | print(logging_config) 28 | 29 | # Enable model invocation logging 30 | try: 31 | bedrock_client.put_model_invocation_logging_configuration( 32 | loggingConfig=logging_config 33 | ) 34 | print('Model invocation logging enabled successfully.') 35 | except Exception as e: 36 | print(f'Error enabling model invocation logging: {e}') 37 | 38 | return log_group 39 | 40 | @helper.delete 41 | def do_delete(event, _): 42 | # Disable model invocation logging 43 | try: 44 | bedrock_client.delete_model_invocation_logging_configuration() 45 | print('Model invocation logging disabled successfully.') 46 | except Exception as e: 47 | print(f'Error disabling model invocation logging: {e}') 48 | 49 | def handler(event, context): 50 | helper(event, context) 51 | 52 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/bedrock-custom/requirements.txt: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | 4 | crhelper -------------------------------------------------------------------------------- /cdk/lib/tenant-template/bedrock.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import * as path from "path"; 5 | import { PythonFunction } from "@aws-cdk/aws-lambda-python-alpha"; 6 | import { Duration } from "aws-cdk-lib"; 7 | import { 8 | ManagedPolicy, 9 | Role, 10 | ServicePrincipal, 11 | PolicyStatement, 12 | Effect, 13 | ArnPrincipal, 14 | } from "aws-cdk-lib/aws-iam"; 15 | import { Function, Runtime } from "aws-cdk-lib/aws-lambda"; 16 | import { LogGroup, RetentionDays } from "aws-cdk-lib/aws-logs"; 17 | import { CustomResource } from "aws-cdk-lib"; 18 | import { Construct } from "constructs"; 19 | 20 | export class BedrockCustom extends Construct { 21 | readonly bedrockLambdaFunc: Function; 22 | readonly modelInvocationLogGroupName: string; 23 | 24 | constructor(scope: Construct, id: string) { 25 | super(scope, id); 26 | 27 | /** 28 | * Creates an IAM role for the Bedrock Customization Lambda function. 29 | * The role is granted read and write access to the Tenant Details table, 30 | * and the ability to put events to the Event Manager. 31 | * The role is also assigned the AWSLambdaBasicExecutionRole, 32 | * CloudWatchLambdaInsightsExecutionRolePolicy, and AWSXrayWriteOnlyAccess 33 | * managed policies. 34 | */ 35 | const bedrockLogGroup = new LogGroup( 36 | this, 37 | "SaaSGenAIWorkshopBedrockLogGroup", 38 | { 39 | retention: RetentionDays.ONE_WEEK, 40 | } 41 | ); 42 | this.modelInvocationLogGroupName = bedrockLogGroup.logGroupName; 43 | 44 | const bedrockLambdaExecRole = new Role(this, "bedrockLambdaExecRole", { 45 | assumedBy: new ServicePrincipal("lambda.amazonaws.com"), 46 | }); 47 | 48 | bedrockLambdaExecRole.addManagedPolicy( 49 | ManagedPolicy.fromAwsManagedPolicyName( 50 | "service-role/AWSLambdaBasicExecutionRole" 51 | ) 52 | ); 53 | bedrockLambdaExecRole.addManagedPolicy( 54 | ManagedPolicy.fromAwsManagedPolicyName( 55 | "CloudWatchLambdaInsightsExecutionRolePolicy" 56 | ) 57 | ); 58 | bedrockLambdaExecRole.addManagedPolicy( 59 | ManagedPolicy.fromAwsManagedPolicyName("AWSXrayWriteOnlyAccess") 60 | ); 61 | 62 | bedrockLambdaExecRole.addToPolicy( 63 | new PolicyStatement({ 64 | actions: [ 65 | "bedrock:PutModelInvocationLoggingConfiguration", 66 | "bedrock:DeleteModelInvocationLoggingConfiguration", 67 | ], 68 | effect: Effect.ALLOW, 69 | resources: ["*"], 70 | }) 71 | ); 72 | 73 | const bedrockLogRole = new Role(this, "bedrockLogRole", { 74 | assumedBy: new ServicePrincipal("bedrock.amazonaws.com"), 75 | }); 76 | 77 | bedrockLogRole.addToPolicy( 78 | new PolicyStatement({ 79 | actions: ["logs:CreateLogStream", "logs:PutLogEvents"], 80 | effect: Effect.ALLOW, 81 | resources: [bedrockLogGroup.logGroupArn], 82 | }) 83 | ); 84 | 85 | bedrockLambdaExecRole.addToPolicy( 86 | new PolicyStatement({ 87 | effect: Effect.ALLOW, 88 | actions: ["iam:PassRole"], 89 | resources: [bedrockLogRole.roleArn], 90 | }) 91 | ); 92 | 93 | /** 94 | * Creates the Bedrock Customization Lambda function. 95 | * The function is configured with the necessary environment variables, 96 | * the Bedrock Customization execution role, and the AWS Lambda Powertools layer. 97 | */ 98 | const bedrockLambdaFunc = new PythonFunction( 99 | this, 100 | "BedrockLambdaFunction", 101 | { 102 | entry: path.join(__dirname, "./bedrock-custom"), 103 | runtime: Runtime.PYTHON_3_12, 104 | index: "bedrock_logs.py", 105 | handler: "handler", 106 | timeout: Duration.seconds(60), 107 | role: bedrockLambdaExecRole, 108 | environment: { 109 | LOG_LEVEL: "INFO", 110 | LOG_GROUP_NAME: bedrockLogGroup.logGroupName, 111 | BEDROCK_LOG_ROLE: bedrockLogRole.roleArn, 112 | }, 113 | } 114 | ); 115 | 116 | // this.bedrockLambdaFunc = bedrockLambdaFunc; 117 | 118 | new CustomResource(scope, `bedrockLogsCustomResource-${id}`, { 119 | serviceToken: bedrockLambdaFunc.functionArn, 120 | properties: { 121 | updateToken: Date.now().toString(), // This will force an update on each deployment 122 | }, 123 | }); 124 | } 125 | } 126 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/bootstrap-template-stack.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | import { 4 | RemovalPolicy, 5 | Stack, 6 | StackProps, 7 | CfnOutput, 8 | Aspects, 9 | } from "aws-cdk-lib"; 10 | import { Construct } from "constructs"; 11 | import { Role, ServicePrincipal, ManagedPolicy } from "aws-cdk-lib/aws-iam"; 12 | import * as path from "path"; 13 | import * as python from "@aws-cdk/aws-lambda-python-alpha"; 14 | import * as lambda from "aws-cdk-lib/aws-lambda"; 15 | 16 | import { IdentityProvider } from "./identity-provider"; 17 | import { opensearchserverless } from "@cdklabs/generative-ai-cdk-constructs"; 18 | import * as aoss from "aws-cdk-lib/aws-opensearchserverless"; 19 | import { Bucket } from "aws-cdk-lib/aws-s3"; 20 | import { ApiGateway } from "./api-gateway"; 21 | import { Services } from "./services"; 22 | import { BedrockCustom } from "./bedrock"; 23 | import { CoreUitlsTemplateStack } from "../core-utils-template-stack"; 24 | import { CostPerTenant } from "./cost-per-tenant"; 25 | import { TenantTokenUsage } from "./tenant-token-usage"; 26 | import { DestroyPolicySetter } from "../destory-policy-setter"; 27 | import { CurAthena } from "./cur-athena"; 28 | import { CostUsageReportUpload } from "./cur-report-upload"; 29 | 30 | interface BootstrapTemplateStackProps extends StackProps { 31 | readonly coreUtilsStack: CoreUitlsTemplateStack; 32 | readonly controlPlaneApiGwUrl: string; 33 | } 34 | 35 | export class BootstrapTemplateStack extends Stack { 36 | constructor( 37 | scope: Construct, 38 | id: string, 39 | props: BootstrapTemplateStackProps 40 | ) { 41 | super(scope, id, props); 42 | 43 | const curPrefix = "CostUsageReport"; 44 | const curDatabaseName = "costexplorerdb"; 45 | // ***************** 46 | // Layers 47 | // ***************** 48 | 49 | // https://docs.powertools.aws.dev/lambda/python/2.31.0/#lambda-layer 50 | const lambdaPowerToolsLayerARN = `arn:aws:lambda:${ 51 | Stack.of(this).region 52 | }:017000801446:layer:AWSLambdaPowertoolsPythonV2:59`; 53 | const lambdaPowerToolsLayer = lambda.LayerVersion.fromLayerVersionArn( 54 | this, 55 | "LambdaPowerTools", 56 | lambdaPowerToolsLayerARN 57 | ); 58 | 59 | const utilsLayer = new python.PythonLayerVersion(this, "UtilsLayer", { 60 | entry: path.join(__dirname, "services/layers/"), 61 | compatibleRuntimes: [lambda.Runtime.PYTHON_3_12], 62 | }); 63 | 64 | const identityProvider = new IdentityProvider(this, "IdentityProvider"); 65 | const app_client_id = 66 | identityProvider.tenantUserPoolClient.userPoolClientId; 67 | const userPoolID = identityProvider.tenantUserPool.userPoolId; 68 | 69 | const tenantTokenUsage = new TenantTokenUsage(this, "TenantTokenUsage", { 70 | lambdaPowerToolsLayer: lambdaPowerToolsLayer, 71 | }); 72 | 73 | // TODO: Lab1 - Add pooled resources 74 | 75 | const collection = new opensearchserverless.VectorCollection( 76 | this, 77 | "SaaSGenAIWorkshopVectorCollection" 78 | ); 79 | 80 | const s3Bucket = new Bucket(this, "SaaSGenAIWorkshopBucket", { 81 | autoDeleteObjects: true, 82 | removalPolicy: RemovalPolicy.DESTROY, 83 | eventBridgeEnabled: true, 84 | }); 85 | 86 | const api = new ApiGateway(this, "SaaSGenAIWorkshopRestApi", {}); 87 | 88 | const services = new Services(this, "SaaSGenAIWorkshopServices", { 89 | appClientID: app_client_id, 90 | userPoolID: userPoolID, 91 | s3Bucket: s3Bucket, 92 | tenantTokenUsageTable: tenantTokenUsage.tenantTokenUsageTable, 93 | restApi: api.restApi, 94 | controlPlaneApiGwUrl: props.controlPlaneApiGwUrl, 95 | lambdaPowerToolsLayer: lambdaPowerToolsLayer, 96 | utilsLayer: utilsLayer, 97 | }); 98 | 99 | const coreUtilsStack = props.coreUtilsStack; 100 | 101 | // Access the codeBuildProject instance from the coreUtilsStack 102 | const codeBuildProject = coreUtilsStack.codeBuildProject; 103 | 104 | const dataAccessPolicy = new aoss.CfnAccessPolicy( 105 | this, 106 | "dataAccessPolicy", 107 | { 108 | name: `${collection.collectionName}`, 109 | description: `Data access policy for: ${collection.collectionName}`, 110 | type: "data", 111 | policy: JSON.stringify([ 112 | { 113 | Rules: [ 114 | { 115 | Resource: [`index/${collection.collectionName}/*`], 116 | Permission: [ 117 | "aoss:CreateIndex", 118 | "aoss:DeleteIndex", 119 | "aoss:UpdateIndex", 120 | "aoss:DescribeIndex", 121 | "aoss:ReadDocument", 122 | "aoss:WriteDocument", 123 | ], 124 | ResourceType: "index", 125 | }, 126 | ], 127 | Principal: [ 128 | // Manual CodeBuild IAM Role implementation until SBT components provide this role. 129 | // As a temporary fix(unitl we have SBT fix), we are adding a dummy role so that we can deploy this data access policy. 130 | // Once workshop is deployed, you need to get CodeBuild project IAM role, add it here and then redeploy. 131 | // temporaryFixRole.roleArn, 132 | codeBuildProject.role?.roleArn, 133 | ], 134 | Description: "saas-genai-data-access-rule", 135 | }, 136 | ]), 137 | } 138 | ); 139 | 140 | new CfnOutput(this, "SaaSGenAIWorkshopOSSDataAccessPolicy", { 141 | value: dataAccessPolicy.name, 142 | }); 143 | 144 | const curS3Bucket = new Bucket(this, "SaaSGenAICURWorkshopBucket", { 145 | autoDeleteObjects: true, 146 | removalPolicy: RemovalPolicy.DESTROY, 147 | }); 148 | 149 | const sourceCodeS3Bucket = new Bucket(this, "TenantSourceCodeBucket", { 150 | autoDeleteObjects: true, 151 | removalPolicy: RemovalPolicy.DESTROY, 152 | }); 153 | 154 | const bedrockCustom = new BedrockCustom(this, "BedrockCustom"); 155 | 156 | const costPerTenant = new CostPerTenant(this, "CostPerTenant", { 157 | lambdaPowerToolsLayer: lambdaPowerToolsLayer, 158 | utilsLayer: utilsLayer, 159 | modelInvocationLogGroupName: bedrockCustom.modelInvocationLogGroupName, 160 | curDatabaseName: curDatabaseName, 161 | tableName: curPrefix.toLowerCase(), 162 | athenaOutputBucketName: curS3Bucket.bucketName, 163 | }); 164 | 165 | const costUsageReportUpload = new CostUsageReportUpload( 166 | this, 167 | "CostUsageReportUpload", 168 | { 169 | curBucketName: curS3Bucket.bucketName, 170 | folderName: curPrefix, 171 | } 172 | ); 173 | 174 | const curAthena = new CurAthena(this, "CurAthena", { 175 | curBucketName: curS3Bucket.bucketName, 176 | folderName: curPrefix, 177 | databaseName: curDatabaseName, 178 | }); 179 | 180 | curAthena.node.addDependency(costUsageReportUpload); 181 | 182 | new CfnOutput(this, "TenantUserpoolId", { 183 | value: identityProvider.tenantUserPool.userPoolId, 184 | }); 185 | 186 | new CfnOutput(this, "UserPoolClientId", { 187 | value: identityProvider.tenantUserPoolClient.userPoolClientId, 188 | }); 189 | 190 | new CfnOutput(this, "ApiGatewayUrl", { 191 | value: api.restApi.url, 192 | }); 193 | 194 | new CfnOutput(this, "ApiGatewayUsagePlan", { 195 | value: api.usagePlanBasicTier.usagePlanId, 196 | }); 197 | 198 | new CfnOutput(this, "SaaSGenAIWorkshopS3Bucket", { 199 | value: s3Bucket.bucketName, 200 | }); 201 | new CfnOutput(this, "SaaSGenAIWorkshopOSSCollectionArn", { 202 | value: collection.collectionArn, 203 | }); 204 | 205 | new CfnOutput(this, "SaaSGenAIWorkshopTriggerIngestionLambdaArn", { 206 | value: services.triggerDataIngestionService.functionArn, 207 | }); 208 | 209 | new CfnOutput(this, "TenantSourceCodeS3Bucket", { 210 | value: sourceCodeS3Bucket.bucketName, 211 | }); 212 | 213 | new CfnOutput(this, "BedrockModelInvocationLogGroupName", { 214 | value: bedrockCustom.modelInvocationLogGroupName, 215 | }); 216 | 217 | Aspects.of(this).add(new DestroyPolicySetter()); 218 | } 219 | } 220 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/cost-per-tenant.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import * as cdk from "aws-cdk-lib"; 5 | import { Construct } from "constructs"; 6 | import { PythonFunction } from "@aws-cdk/aws-lambda-python-alpha"; 7 | import * as iam from "aws-cdk-lib/aws-iam"; 8 | import * as lambda from "aws-cdk-lib/aws-lambda"; 9 | import * as dynamodb from "aws-cdk-lib/aws-dynamodb"; 10 | import * as events from "aws-cdk-lib/aws-events"; 11 | import * as targets from "aws-cdk-lib/aws-events-targets"; 12 | 13 | import * as path from "path"; 14 | 15 | export interface CostPerTenantProps { 16 | readonly lambdaPowerToolsLayer: lambda.ILayerVersion; 17 | readonly utilsLayer: lambda.ILayerVersion; 18 | readonly modelInvocationLogGroupName: string; 19 | readonly curDatabaseName: string; 20 | readonly tableName: string; 21 | readonly athenaOutputBucketName: string; 22 | } 23 | 24 | export class CostPerTenant extends Construct { 25 | constructor(scope: Construct, id: string, props: CostPerTenantProps) { 26 | super(scope, id); 27 | 28 | const region = cdk.Stack.of(this).region; 29 | const accountId = cdk.Stack.of(this).account; 30 | const athenaOutputPath = `${props.athenaOutputBucketName}/athenaoutput`; 31 | 32 | // TODO: Athena setup 33 | 34 | const partitionKey = { name: "Date", type: dynamodb.AttributeType.NUMBER }; 35 | const sortKey = { 36 | name: "TenantId#ServiceName", 37 | type: dynamodb.AttributeType.STRING, 38 | }; 39 | 40 | // Create the DynamoDB table 41 | const table = new dynamodb.TableV2(this, "TenantCostAndUsageAttribution", { 42 | tableName: "TenantCostAndUsageAttribution", 43 | partitionKey: partitionKey, 44 | sortKey: sortKey, 45 | }); 46 | 47 | const tenantCostCalculatorLambdaExecRole = new iam.Role( 48 | this, 49 | "tenantCostCalculatorLambdaExecRole", 50 | { 51 | assumedBy: new iam.ServicePrincipal("lambda.amazonaws.com"), 52 | managedPolicies: [ 53 | iam.ManagedPolicy.fromAwsManagedPolicyName( 54 | "CloudWatchLambdaInsightsExecutionRolePolicy" 55 | ), 56 | iam.ManagedPolicy.fromAwsManagedPolicyName( 57 | "service-role/AWSLambdaBasicExecutionRole" 58 | ), 59 | ], 60 | } 61 | ); 62 | 63 | tenantCostCalculatorLambdaExecRole.addToPolicy( 64 | new iam.PolicyStatement({ 65 | effect: iam.Effect.ALLOW, 66 | actions: ["dynamodb:PutItem"], 67 | resources: [table.tableArn], 68 | }) 69 | ); 70 | 71 | tenantCostCalculatorLambdaExecRole.addToPolicy( 72 | new iam.PolicyStatement({ 73 | effect: iam.Effect.ALLOW, 74 | actions: [ 75 | "logs:GetQueryResults", 76 | "logs:StartQuery", 77 | "logs:StopQuery", 78 | "logs:FilterLogEvents", 79 | "logs:DescribeLogGroups", 80 | ], 81 | resources: [`arn:aws:logs:${region}:${accountId}:log-group:*`], 82 | }) 83 | ); 84 | 85 | tenantCostCalculatorLambdaExecRole.addToPolicy( 86 | new iam.PolicyStatement({ 87 | effect: iam.Effect.ALLOW, 88 | actions: ["cloudformation:ListStackResources"], 89 | resources: [`arn:aws:cloudformation:${region}:${accountId}:stack/*/*`], 90 | }) 91 | ); 92 | 93 | tenantCostCalculatorLambdaExecRole.addToPolicy( 94 | new iam.PolicyStatement({ 95 | effect: iam.Effect.ALLOW, 96 | actions: [ 97 | "athena:StartQueryExecution", 98 | "athena:StopQueryExecution", 99 | "athena:GetQueryExecution", 100 | "athena:GetQueryResults", 101 | ], 102 | resources: [`arn:aws:athena:${region}:${accountId}:workgroup/*`], 103 | }) 104 | ); 105 | 106 | tenantCostCalculatorLambdaExecRole.addToPolicy( 107 | new iam.PolicyStatement({ 108 | effect: iam.Effect.ALLOW, 109 | actions: [ 110 | "s3:PutObject", 111 | "s3:AbortMultipartUpload", 112 | "s3:ListMultipartUploadParts", 113 | "s3:ListBucketMultipartUploads", 114 | "s3:GetObject", 115 | "s3:ListBucket", 116 | "s3:GetBucketLocation", 117 | ], 118 | resources: [ 119 | `arn:aws:s3:::${props.athenaOutputBucketName}`, 120 | `arn:aws:s3:::${props.athenaOutputBucketName}/*`, 121 | ], 122 | }) 123 | ); 124 | 125 | tenantCostCalculatorLambdaExecRole.addToPolicy( 126 | new iam.PolicyStatement({ 127 | effect: iam.Effect.ALLOW, 128 | actions: ["glue:GetDatabase", "glue:GetTable", "glue:GetPartitions"], 129 | resources: [`*`], 130 | }) 131 | ); 132 | 133 | const tenantCostCalculatorService = new PythonFunction( 134 | this, 135 | "TenantCostCalculatorService", 136 | { 137 | functionName: "TenantCostCalculatorService", 138 | entry: path.join(__dirname, "services/aggregate-metrics/"), 139 | runtime: lambda.Runtime.PYTHON_3_12, 140 | index: "tenant_cost_calculator.py", 141 | handler: "calculate_cost_per_tenant", 142 | timeout: cdk.Duration.seconds(60), 143 | role: tenantCostCalculatorLambdaExecRole, 144 | layers: [props.lambdaPowerToolsLayer, props.utilsLayer], 145 | environment: { 146 | ATHENA_S3_OUTPUT: athenaOutputPath, 147 | TENANT_COST_DYNAMODB_TABLE: table.tableName, 148 | MODEL_INVOCATION_LOG_GROUPNAME: props.modelInvocationLogGroupName, 149 | CUR_DATABASE_NAME: props.curDatabaseName, 150 | CUR_TABLE_NAME: props.tableName, 151 | }, 152 | } 153 | ); 154 | 155 | const rule = new events.Rule(this, "ScheduleRule", { 156 | schedule: events.Schedule.rate(cdk.Duration.minutes(5)), 157 | }); 158 | 159 | // Add the Lambda function as a target of the rule 160 | rule.addTarget(new targets.LambdaFunction(tenantCostCalculatorService)); 161 | } 162 | } 163 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/cur-athena.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import * as cdk from "aws-cdk-lib"; 5 | import * as glue from "aws-cdk-lib/aws-glue"; 6 | import * as iam from "aws-cdk-lib/aws-iam"; 7 | import * as lambda from "aws-cdk-lib/aws-lambda"; 8 | import * as cr from "aws-cdk-lib/custom-resources"; 9 | import { Construct } from "constructs"; 10 | 11 | export interface CurAthenaProps extends cdk.StackProps { 12 | curBucketName: string; 13 | folderName: string; 14 | databaseName: string; 15 | } 16 | 17 | export class CurAthena extends Construct { 18 | public constructor(scope: Construct, id: string, props: CurAthenaProps) { 19 | super(scope, id); 20 | 21 | const curBucketName = props.curBucketName; 22 | const folderName = props.folderName; 23 | const accountId = cdk.Stack.of(this).account; 24 | const databaseName = props.databaseName; 25 | 26 | // Resources 27 | const awscurCrawlerComponentFunction = new iam.CfnRole( 28 | this, 29 | "AWSCURCrawlerComponentFunction", 30 | { 31 | assumeRolePolicyDocument: { 32 | Version: "2012-10-17", 33 | Statement: [ 34 | { 35 | Effect: "Allow", 36 | Principal: { 37 | Service: ["glue.amazonaws.com"], 38 | }, 39 | Action: ["sts:AssumeRole"], 40 | }, 41 | ], 42 | }, 43 | path: "/", 44 | managedPolicyArns: [ 45 | `arn:aws:iam::aws:policy/service-role/AWSGlueServiceRole`, 46 | ], 47 | policies: [ 48 | { 49 | policyName: "AWSCURCrawlerComponentFunction", 50 | policyDocument: { 51 | Version: "2012-10-17", 52 | Statement: [ 53 | { 54 | Effect: "Allow", 55 | Action: [ 56 | "logs:CreateLogGroup", 57 | "logs:CreateLogStream", 58 | "logs:PutLogEvents", 59 | ], 60 | Resource: `arn:aws:logs:*:*:*`, 61 | }, 62 | { 63 | Effect: "Allow", 64 | Action: [ 65 | "glue:UpdateDatabase", 66 | "glue:UpdatePartition", 67 | "glue:CreateTable", 68 | "glue:UpdateTable", 69 | "glue:ImportCatalogToGlue", 70 | ], 71 | Resource: "*", 72 | }, 73 | { 74 | Effect: "Allow", 75 | Action: ["s3:GetObject", "s3:PutObject"], 76 | Resource: `arn:aws:s3:::${curBucketName}/${folderName}*`, 77 | }, 78 | ], 79 | }, 80 | }, 81 | { 82 | policyName: "AWSCURKMSDecryption", 83 | policyDocument: { 84 | Version: "2012-10-17", 85 | Statement: [ 86 | { 87 | Effect: "Allow", 88 | Action: ["kms:Decrypt"], 89 | Resource: "*", 90 | }, 91 | ], 92 | }, 93 | }, 94 | ], 95 | } 96 | ); 97 | 98 | const awscurCrawlerLambdaExecutor = new iam.CfnRole( 99 | this, 100 | "AWSCURCrawlerLambdaExecutor", 101 | { 102 | assumeRolePolicyDocument: { 103 | Version: "2012-10-17", 104 | Statement: [ 105 | { 106 | Effect: "Allow", 107 | Principal: { 108 | Service: ["lambda.amazonaws.com"], 109 | }, 110 | Action: ["sts:AssumeRole"], 111 | }, 112 | ], 113 | }, 114 | path: "/", 115 | policies: [ 116 | { 117 | policyName: "AWSCURCrawlerLambdaExecutor", 118 | policyDocument: { 119 | Version: "2012-10-17", 120 | Statement: [ 121 | { 122 | Effect: "Allow", 123 | Action: [ 124 | "logs:CreateLogGroup", 125 | "logs:CreateLogStream", 126 | "logs:PutLogEvents", 127 | ], 128 | Resource: `arn:aws:logs:*:*:*`, 129 | }, 130 | { 131 | Effect: "Allow", 132 | Action: ["glue:StartCrawler"], 133 | Resource: "*", 134 | }, 135 | ], 136 | }, 137 | }, 138 | ], 139 | } 140 | ); 141 | 142 | const awscurDatabase = new glue.CfnDatabase(this, "AWSCURDatabase", { 143 | databaseInput: { 144 | name: databaseName, 145 | }, 146 | catalogId: accountId, 147 | }); 148 | 149 | const awss3curLambdaExecutor = new iam.CfnRole( 150 | this, 151 | "AWSS3CURLambdaExecutor", 152 | { 153 | assumeRolePolicyDocument: { 154 | Version: "2012-10-17", 155 | Statement: [ 156 | { 157 | Effect: "Allow", 158 | Principal: { 159 | Service: ["lambda.amazonaws.com"], 160 | }, 161 | Action: ["sts:AssumeRole"], 162 | }, 163 | ], 164 | }, 165 | path: "/", 166 | policies: [ 167 | { 168 | policyName: "AWSS3CURLambdaExecutor", 169 | policyDocument: { 170 | Version: "2012-10-17", 171 | Statement: [ 172 | { 173 | Effect: "Allow", 174 | Action: [ 175 | "logs:CreateLogGroup", 176 | "logs:CreateLogStream", 177 | "logs:PutLogEvents", 178 | ], 179 | Resource: `arn:aws:logs:*:*:*`, 180 | }, 181 | { 182 | Effect: "Allow", 183 | Action: ["s3:PutBucketNotification"], 184 | Resource: `arn:aws:s3:::${curBucketName}`, 185 | }, 186 | ], 187 | }, 188 | }, 189 | ], 190 | } 191 | ); 192 | 193 | const awscurCrawler = new glue.CfnCrawler(this, "AWSCURCrawler", { 194 | name: "AWSCURCrawler-" + folderName.toLowerCase(), 195 | description: 196 | "A recurring crawler that keeps your CUR table in Athena up-to-date.", 197 | role: awscurCrawlerComponentFunction.attrArn, 198 | databaseName: awscurDatabase.ref, 199 | targets: { 200 | s3Targets: [ 201 | { 202 | path: "s3://" + curBucketName + "/" + folderName + "/", 203 | exclusions: [ 204 | "**.json", 205 | "**.yml", 206 | "**.sql", 207 | "**.csv", 208 | "**.gz", 209 | "**.zip", 210 | ], 211 | }, 212 | ], 213 | }, 214 | schemaChangePolicy: { 215 | updateBehavior: "UPDATE_IN_DATABASE", 216 | deleteBehavior: "DELETE_FROM_DATABASE", 217 | }, 218 | }); 219 | awscurCrawler.addDependency(awscurDatabase); 220 | awscurCrawler.addDependency(awscurCrawlerComponentFunction); 221 | 222 | const awscurInitializer = new lambda.CfnFunction( 223 | this, 224 | "AWSCURInitializer", 225 | { 226 | code: { 227 | zipFile: 228 | "const { GlueClient, StartCrawlerCommand } = require('@aws-sdk/client-glue'); const response = require('./cfn-response'); exports.handler = function (event, context, callback) {\n if (event.RequestType === 'Delete') {\n response.send(event, context, response.SUCCESS);\n } else {\n const glue = new GlueClient();\n const input = {\n Name: 'AWSCURCrawler-" + 229 | folderName.toLowerCase() + 230 | "',\n };\n const command = new StartCrawlerCommand(input);\n glue.send(command, function (err, data) {\n if (err) {\n const responseData = JSON.parse(this.httpResponse.body);\n if (responseData['__type'] == 'CrawlerRunningException') {\n callback(null, responseData.Message);\n } else {\n const responseString = JSON.stringify(responseData);\n if (event.ResponseURL) {\n response.send(event, context, response.FAILED, {\n msg: responseString,\n });\n } else {\n callback(responseString);\n }\n }\n } else {\n if (event.ResponseURL) {\n response.send(event, context, response.SUCCESS);\n } else {\n callback(null, response.SUCCESS);\n }\n }\n });\n }\n};\n", 231 | }, 232 | handler: "index.handler", 233 | timeout: 30, 234 | runtime: "nodejs18.x", 235 | reservedConcurrentExecutions: 1, 236 | role: awscurCrawlerLambdaExecutor.attrArn, 237 | } 238 | ); 239 | awscurInitializer.addDependency(awscurCrawler); 240 | 241 | const awss3curEventLambdaPermission = new lambda.CfnPermission( 242 | this, 243 | "AWSS3CUREventLambdaPermission", 244 | { 245 | action: "lambda:InvokeFunction", 246 | functionName: awscurInitializer.attrArn, 247 | principal: "s3.amazonaws.com", 248 | sourceAccount: accountId, 249 | sourceArn: `arn:aws:s3:::${curBucketName}`, 250 | } 251 | ); 252 | 253 | const awss3curNotification = new lambda.CfnFunction( 254 | this, 255 | "AWSS3CURNotification", 256 | { 257 | code: { 258 | zipFile: 259 | "const { S3Client, PutBucketNotificationConfigurationCommand } = require('@aws-sdk/client-s3'); const response = require('./cfn-response'); exports.handler = function (event, context, callback) {\n const s3 = new S3Client();\n const putConfigRequest = function (notificationConfiguration) {\n return new Promise(function (resolve, reject) {\n const input = {\n Bucket: event.ResourceProperties.BucketName,\n NotificationConfiguration: notificationConfiguration,\n };\n const command = new PutBucketNotificationConfigurationCommand(input);\n s3.send(command, function (err, data) {\n if (err)\n reject({\n msg: this.httpResponse.body.toString(),\n error: err,\n data: data,\n });\n else resolve(data);\n });\n });\n };\n const newNotificationConfig = {};\n if (event.RequestType !== 'Delete') {\n newNotificationConfig.LambdaFunctionConfigurations = [\n {\n Events: ['s3:ObjectCreated:*'],\n LambdaFunctionArn:\n event.ResourceProperties.TargetLambdaArn || 'missing arn',\n Filter: {\n Key: {\n FilterRules: [\n { Name: 'prefix', Value: event.ResourceProperties.ReportKey },\n ],\n },\n },\n },\n ];\n }\n putConfigRequest(newNotificationConfig)\n .then(function (result) {\n response.send(event, context, response.SUCCESS, result);\n callback(null, result);\n })\n .catch(function (error) {\n response.send(event, context, response.FAILED, error);\n console.log(error);\n callback(error);\n });\n};\n", 260 | }, 261 | handler: "index.handler", 262 | timeout: 30, 263 | runtime: "nodejs18.x", 264 | reservedConcurrentExecutions: 1, 265 | role: awss3curLambdaExecutor.attrArn, 266 | } 267 | ); 268 | awss3curNotification.addDependency(awscurInitializer); 269 | awss3curNotification.addDependency(awss3curEventLambdaPermission); 270 | awss3curNotification.addDependency(awss3curLambdaExecutor); 271 | 272 | const awsStartCURCrawler = new cdk.CfnCustomResource( 273 | this, 274 | "AWSStartCURCrawler", 275 | { 276 | serviceToken: awscurInitializer.attrArn, 277 | } 278 | ); 279 | awsStartCURCrawler.addDependency(awscurInitializer); 280 | 281 | const awsPutS3CURNotification = new cdk.CustomResource( 282 | this, 283 | "AWSPutS3CURNotification", 284 | { 285 | serviceToken: awss3curNotification.attrArn, 286 | properties: { 287 | TargetLambdaArn: awscurInitializer.attrArn, 288 | BucketName: curBucketName, 289 | ReportKey: `${folderName}`, 290 | }, 291 | } 292 | ); 293 | awsPutS3CURNotification.node.addDependency(awss3curNotification); 294 | awsPutS3CURNotification.node.addDependency(awscurInitializer); 295 | } 296 | } 297 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/cur-report-upload.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import * as cdk from "aws-cdk-lib"; 5 | import * as s3 from "aws-cdk-lib/aws-s3"; 6 | import * as s3deploy from "aws-cdk-lib/aws-s3-deployment"; 7 | import { Bucket } from "aws-cdk-lib/aws-s3"; 8 | import { Construct } from "constructs"; 9 | 10 | export interface CostUsageReportUploadProps extends cdk.StackProps { 11 | curBucketName: string; 12 | folderName: string; 13 | } 14 | 15 | export class CostUsageReportUpload extends Construct { 16 | constructor(scope: Construct, id: string, props: CostUsageReportUploadProps) { 17 | super(scope, id); 18 | 19 | const curBucketName = s3.Bucket.fromBucketName( 20 | this, 21 | "curBucketName", 22 | props.curBucketName 23 | ); 24 | const folderName = props.folderName; 25 | 26 | // Deploy the local folder to the S3 bucket 27 | new s3deploy.BucketDeployment(this, "DeployCostUsageReports", { 28 | sources: [s3deploy.Source.asset("../data/cur_report.zip")], 29 | destinationBucket: curBucketName, 30 | destinationKeyPrefix: folderName, 31 | }); 32 | } 33 | } 34 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/identity-provider.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import { aws_cognito, StackProps, Duration } from "aws-cdk-lib"; 5 | import { Construct } from "constructs"; 6 | import { IdentityDetails } from "../interfaces/identity-details"; 7 | import { KnowledgeBase } from "@cdklabs/generative-ai-cdk-constructs/lib/cdk-lib/bedrock"; 8 | 9 | export class IdentityProvider extends Construct { 10 | public readonly tenantUserPool: aws_cognito.UserPool; 11 | public readonly tenantUserPoolClient: aws_cognito.UserPoolClient; 12 | public readonly identityDetails: IdentityDetails; 13 | constructor(scope: Construct, id: string, props?: StackProps) { 14 | super(scope, id); 15 | 16 | this.tenantUserPool = new aws_cognito.UserPool(this, "tenantUserPool", { 17 | autoVerify: { email: true }, 18 | accountRecovery: aws_cognito.AccountRecovery.EMAIL_ONLY, 19 | standardAttributes: { 20 | email: { 21 | required: true, 22 | mutable: true, 23 | }, 24 | }, 25 | customAttributes: { 26 | tenantId: new aws_cognito.StringAttribute({ 27 | mutable: true, 28 | }), 29 | userRole: new aws_cognito.StringAttribute({ 30 | mutable: true, 31 | }), 32 | }, 33 | }); 34 | 35 | const writeAttributes = new aws_cognito.ClientAttributes() 36 | .withStandardAttributes({ email: true }) 37 | .withCustomAttributes("tenantId", "userRole"); 38 | 39 | this.tenantUserPoolClient = new aws_cognito.UserPoolClient( 40 | this, 41 | "tenantUserPoolClient", 42 | { 43 | userPool: this.tenantUserPool, 44 | generateSecret: false, 45 | accessTokenValidity: Duration.minutes(180), 46 | idTokenValidity: Duration.minutes(180), 47 | authFlows: { 48 | userPassword: true, 49 | adminUserPassword: false, 50 | userSrp: true, 51 | custom: false, 52 | }, 53 | writeAttributes: writeAttributes, 54 | oAuth: { 55 | scopes: [ 56 | aws_cognito.OAuthScope.EMAIL, 57 | aws_cognito.OAuthScope.OPENID, 58 | aws_cognito.OAuthScope.PROFILE, 59 | ], 60 | flows: { 61 | authorizationCodeGrant: true, 62 | implicitCodeGrant: true, 63 | }, 64 | }, 65 | } 66 | ); 67 | 68 | this.identityDetails = { 69 | name: "Cognito", 70 | details: { 71 | userPoolId: this.tenantUserPool.userPoolId, 72 | appClientId: this.tenantUserPoolClient.userPoolClientId, 73 | }, 74 | }; 75 | } 76 | } 77 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import { Construct } from "constructs"; 5 | import * as path from "path"; 6 | import * as python from "@aws-cdk/aws-lambda-python-alpha"; 7 | import { 8 | Architecture, 9 | Code, 10 | Runtime, 11 | LayerVersion, 12 | Function, 13 | ILayerVersion, 14 | } from "aws-cdk-lib/aws-lambda"; 15 | import { Duration, Stack, Arn } from "aws-cdk-lib"; 16 | import { Bucket } from "aws-cdk-lib/aws-s3"; 17 | import { 18 | Role, 19 | ServicePrincipal, 20 | PolicyStatement, 21 | Effect, 22 | ArnPrincipal, 23 | ManagedPolicy, 24 | PolicyDocument, 25 | } from "aws-cdk-lib/aws-iam"; 26 | import { 27 | RestApi, 28 | LambdaIntegration, 29 | AuthorizationType, 30 | } from "aws-cdk-lib/aws-apigateway"; 31 | import * as apigw from "aws-cdk-lib/aws-apigateway"; 32 | import { Asset } from "aws-cdk-lib/aws-s3-assets"; 33 | import { TableV2 } from "aws-cdk-lib/aws-dynamodb"; 34 | 35 | export interface ServicesProps { 36 | readonly appClientID: string; 37 | readonly userPoolID: string; 38 | readonly s3Bucket: Bucket; 39 | readonly tenantTokenUsageTable: TableV2; 40 | readonly restApi: RestApi; 41 | readonly controlPlaneApiGwUrl: string; 42 | readonly lambdaPowerToolsLayer: ILayerVersion; 43 | readonly utilsLayer: ILayerVersion; 44 | } 45 | 46 | export class Services extends Construct { 47 | public readonly ragService: Function; 48 | public readonly s3UploaderService: Function; 49 | public readonly triggerDataIngestionService: Function; 50 | public readonly getJWTTokenService: Function; 51 | public readonly authorizerService: Function; 52 | 53 | constructor(scope: Construct, id: string, props: ServicesProps) { 54 | super(scope, id); 55 | 56 | const region = Stack.of(this).region; 57 | const accountId = Stack.of(this).account; 58 | 59 | const invoke = props.restApi.root.addResource("invoke"); 60 | const s3Upload = props.restApi.root.addResource("upload"); 61 | 62 | // ***************** 63 | // Authorizer Lambda 64 | // ***************** 65 | 66 | const authorizerLambdaExecRole = new Role( 67 | this, 68 | "authorizerLambdaExecRole", 69 | { 70 | assumedBy: new ServicePrincipal("lambda.amazonaws.com"), 71 | managedPolicies: [ 72 | ManagedPolicy.fromAwsManagedPolicyName( 73 | "CloudWatchLambdaInsightsExecutionRolePolicy" 74 | ), 75 | ManagedPolicy.fromAwsManagedPolicyName( 76 | "service-role/AWSLambdaBasicExecutionRole" 77 | ), 78 | ], 79 | } 80 | ); 81 | 82 | authorizerLambdaExecRole.addToPolicy( 83 | new PolicyStatement({ 84 | effect: Effect.ALLOW, 85 | actions: ["cognito-idp:InitiateAuth"], 86 | resources: [ 87 | `arn:aws:cognito-idp:${region}:${accountId}:userpool/${props.userPoolID}`, 88 | ], 89 | }) 90 | ); 91 | 92 | const tenantTokenUsageTableAccessRole = new Role( 93 | this, 94 | "TenantTokenUsageTableAccessRole", 95 | { 96 | assumedBy: new ServicePrincipal("lambda.amazonaws.com"), 97 | inlinePolicies: { 98 | DynamoDBPolicy: new PolicyDocument({ 99 | statements: [ 100 | new PolicyStatement({ 101 | effect: Effect.ALLOW, 102 | actions: ["dynamodb:GetItem"], 103 | resources: [props.tenantTokenUsageTable.tableArn], 104 | conditions: { 105 | "ForAllValues:StringEquals": { 106 | "dynamodb:LeadingKeys": ["${aws:PrincipalTag/TenantId}"], 107 | }, 108 | }, 109 | }), 110 | ], 111 | }), 112 | }, 113 | } 114 | ); 115 | 116 | tenantTokenUsageTableAccessRole.assumeRolePolicy?.addStatements( 117 | new PolicyStatement({ 118 | actions: ["sts:AssumeRole", "sts:TagSession"], 119 | effect: Effect.ALLOW, 120 | principals: [new ArnPrincipal(authorizerLambdaExecRole.roleArn)], 121 | conditions: { 122 | StringLike: { 123 | "aws:RequestTag/TenantId": "*", 124 | }, 125 | }, 126 | }) 127 | ); 128 | 129 | // ********************* 130 | // Combined ABAC Role 131 | // ********************* 132 | 133 | const abacExecRole = new Role(this, "AbacExecRole", { 134 | assumedBy: new ServicePrincipal("lambda.amazonaws.com"), 135 | managedPolicies: [ 136 | ManagedPolicy.fromAwsManagedPolicyName( 137 | "CloudWatchLambdaInsightsExecutionRolePolicy" 138 | ), 139 | ManagedPolicy.fromAwsManagedPolicyName( 140 | "service-role/AWSLambdaBasicExecutionRole" 141 | ), 142 | ], 143 | inlinePolicies: { 144 | ABACPolicy: new PolicyDocument({ 145 | statements: [ 146 | new PolicyStatement({ 147 | effect: Effect.ALLOW, 148 | actions: ["bedrock:ListKnowledgeBases"], 149 | resources: ["*"], 150 | }), 151 | new PolicyStatement({ 152 | effect: Effect.ALLOW, 153 | actions: ["bedrock:InvokeModel"], 154 | resources: [ 155 | "arn:aws:bedrock:*::foundation-model/amazon.titan-text-lite-v1", 156 | ], 157 | }), 158 | new PolicyStatement({ 159 | effect: Effect.ALLOW, 160 | actions: [ 161 | "s3:PutObject", 162 | "s3:GetObject", 163 | "s3:ListBucket", 164 | "s3:DeleteObject", 165 | ], 166 | resources: [ 167 | `arn:aws:s3:::${props.s3Bucket.bucketName}` + 168 | "/${aws:PrincipalTag/TenantId}/*", 169 | ], 170 | }), 171 | new PolicyStatement({ 172 | effect: Effect.ALLOW, 173 | actions: ["bedrock:Retrieve", "bedrock:RetrieveAndGenerate"], 174 | resources: [ 175 | Arn.format( 176 | { 177 | service: "bedrock", 178 | resource: "knowledge-base", 179 | // TODO: Lab2 - Add principalTag in ABAC policy 180 | resourceName: "*", 181 | account: accountId, 182 | region: region, 183 | }, 184 | Stack.of(this) 185 | ), 186 | ], 187 | }), 188 | ], 189 | }), 190 | }, 191 | }); 192 | 193 | abacExecRole.assumeRolePolicy?.addStatements( 194 | new PolicyStatement({ 195 | actions: ["sts:AssumeRole", "sts:TagSession"], 196 | effect: Effect.ALLOW, 197 | principals: [new ArnPrincipal(authorizerLambdaExecRole.roleArn)], 198 | conditions: { 199 | StringLike: { 200 | "aws:RequestTag/TenantId": "*", 201 | "aws:RequestTag/KnowledgeBaseId": "*", 202 | }, 203 | }, 204 | }) 205 | ); 206 | 207 | const authorizerService = new python.PythonFunction( 208 | this, 209 | "AuthorizerService", 210 | { 211 | functionName: "authorizerService", 212 | entry: path.join(__dirname, "services/authorizerService/"), 213 | runtime: Runtime.PYTHON_3_12, 214 | index: "tenant_authorizer.py", 215 | handler: "lambda_handler", 216 | timeout: Duration.seconds(60), 217 | role: authorizerLambdaExecRole, 218 | layers: [props.lambdaPowerToolsLayer, props.utilsLayer], 219 | environment: { 220 | APP_CLIENT_ID: props.appClientID, 221 | USER_POOL_ID: props.userPoolID, 222 | ASSUME_ROLE_ARN: abacExecRole.roleArn, 223 | CP_API_GW_URL: props.controlPlaneApiGwUrl, 224 | TENANT_TOKEN_USAGE_DYNAMODB_TABLE: 225 | props.tenantTokenUsageTable.tableName, 226 | TENANT_TOKEN_USAGE_ROLE_ARN: tenantTokenUsageTableAccessRole.roleArn, 227 | }, 228 | } 229 | ); 230 | 231 | const authorizer = new apigw.RequestAuthorizer( 232 | this, 233 | "apiRequestAuthorizer", 234 | { 235 | handler: authorizerService, 236 | identitySources: [apigw.IdentitySource.header("authorization")], 237 | resultsCacheTtl: Duration.seconds(0), 238 | } 239 | ); 240 | 241 | // RAG lambda 242 | const raglambdaExecRole = new Role(this, "RaglambdaExecRole", { 243 | assumedBy: new ServicePrincipal("lambda.amazonaws.com"), 244 | managedPolicies: [ 245 | ManagedPolicy.fromAwsManagedPolicyName( 246 | "CloudWatchLambdaInsightsExecutionRolePolicy" 247 | ), 248 | ManagedPolicy.fromAwsManagedPolicyName( 249 | "service-role/AWSLambdaBasicExecutionRole" 250 | ), 251 | ], 252 | }); 253 | 254 | const ragService = new python.PythonFunction(this, "RagService", { 255 | functionName: "ragService", 256 | entry: path.join(__dirname, "services/ragService/"), 257 | runtime: Runtime.PYTHON_3_12, 258 | index: "rag_service.py", 259 | handler: "lambda_handler", 260 | timeout: Duration.seconds(60), 261 | memorySize: 256, 262 | role: raglambdaExecRole, 263 | layers: [props.lambdaPowerToolsLayer, props.utilsLayer], 264 | environment: { 265 | POWERTOOLS_SERVICE_NAME: "RagService", 266 | POWERTOOLS_METRICS_NAMESPACE: "SaaSRAGGenAI", 267 | }, 268 | }); 269 | 270 | this.ragService = ragService; 271 | invoke.addMethod( 272 | "POST", 273 | new LambdaIntegration(this.ragService, { proxy: true }), 274 | { 275 | authorizer: authorizer, 276 | authorizationType: apigw.AuthorizationType.CUSTOM, 277 | apiKeyRequired: true, 278 | } 279 | ); 280 | 281 | // S3 Uploader lambda 282 | const s3UploaderExecRole = new Role(this, "S3UploaderExecRole", { 283 | assumedBy: new ServicePrincipal("lambda.amazonaws.com"), 284 | managedPolicies: [ 285 | ManagedPolicy.fromAwsManagedPolicyName( 286 | "CloudWatchLambdaInsightsExecutionRolePolicy" 287 | ), 288 | ManagedPolicy.fromAwsManagedPolicyName( 289 | "service-role/AWSLambdaBasicExecutionRole" 290 | ), 291 | ], 292 | }); 293 | 294 | const s3Uploader = new python.PythonFunction(this, "S3Uploader", { 295 | functionName: "s3Uploader", 296 | entry: path.join(__dirname, "services/s3Uploader/"), 297 | runtime: Runtime.PYTHON_3_12, 298 | index: "s3uploader.py", 299 | handler: "lambda_handler", 300 | timeout: Duration.seconds(60), 301 | role: s3UploaderExecRole, 302 | layers: [props.lambdaPowerToolsLayer], 303 | environment: { 304 | S3_BUCKET_NAME: props.s3Bucket.bucketName, 305 | }, 306 | }); 307 | 308 | this.s3UploaderService = s3Uploader; 309 | s3Upload.addMethod( 310 | "POST", 311 | new LambdaIntegration(this.s3UploaderService, { proxy: true }), 312 | { 313 | authorizer: authorizer, 314 | authorizationType: apigw.AuthorizationType.CUSTOM, 315 | } 316 | ); 317 | 318 | // Trigger data ingestion lambda 319 | const triggerDataIngestionExecRole = new Role( 320 | this, 321 | "TriggerDataIngestionExecRole", 322 | { 323 | assumedBy: new ServicePrincipal("lambda.amazonaws.com"), 324 | managedPolicies: [ 325 | ManagedPolicy.fromAwsManagedPolicyName( 326 | "CloudWatchLambdaInsightsExecutionRolePolicy" 327 | ), 328 | ManagedPolicy.fromAwsManagedPolicyName( 329 | "service-role/AWSLambdaBasicExecutionRole" 330 | ), 331 | ], 332 | } 333 | ); 334 | 335 | // ABAC role which will be assumed by the Data Ingestion lambda 336 | const triggerDataIngestionServiceAssumeRole = new Role( 337 | this, 338 | "TriggerDataIngestionServiceAssumeRole", 339 | { 340 | assumedBy: new ServicePrincipal("lambda.amazonaws.com"), 341 | } 342 | ); 343 | 344 | triggerDataIngestionServiceAssumeRole.assumeRolePolicy?.addStatements( 345 | new PolicyStatement({ 346 | actions: ["sts:AssumeRole", "sts:TagSession"], 347 | effect: Effect.ALLOW, 348 | principals: [new ArnPrincipal(triggerDataIngestionExecRole.roleArn)], 349 | }) 350 | ); 351 | 352 | triggerDataIngestionServiceAssumeRole.addToPolicy( 353 | new PolicyStatement({ 354 | effect: Effect.ALLOW, 355 | actions: ["bedrock:StartIngestionJob", "bedrock:GetIngestionJob"], 356 | resources: [ 357 | Arn.format( 358 | { 359 | service: "bedrock", 360 | resource: "knowledge-base", 361 | resourceName: "${aws:PrincipalTag/KnowledgeBaseId}", 362 | account: accountId, 363 | region: region, 364 | }, 365 | Stack.of(this) 366 | ), 367 | ], 368 | }) 369 | ); 370 | 371 | const triggerDataIngestionService = new python.PythonFunction( 372 | this, 373 | "TriggerDataIngestionService", 374 | { 375 | functionName: "triggerDataIngestionService", 376 | entry: path.join(__dirname, "services/triggerDataIngestionService/"), 377 | runtime: Runtime.PYTHON_3_12, 378 | index: "trigger_data_ingestion.py", 379 | handler: "lambda_handler", 380 | timeout: Duration.seconds(60), 381 | role: triggerDataIngestionExecRole, 382 | layers: [props.lambdaPowerToolsLayer], 383 | environment: { 384 | ASSUME_ROLE_ARN: triggerDataIngestionServiceAssumeRole.roleArn, 385 | }, 386 | } 387 | ); 388 | 389 | this.triggerDataIngestionService = triggerDataIngestionService; 390 | 391 | // Add permission for eventbrige to trigger data ingestion service 392 | const eventBusRuleArn = Arn.format( 393 | { 394 | service: "events", 395 | resource: "rule/*", 396 | account: accountId, 397 | region: region, 398 | }, 399 | Stack.of(this) 400 | ); 401 | triggerDataIngestionService.addPermission( 402 | "EventBusTriggerDataIngestionPermission", 403 | { 404 | principal: new ServicePrincipal("events.amazonaws.com"), 405 | action: "lambda:InvokeFunction", 406 | sourceArn: eventBusRuleArn, 407 | } 408 | ); 409 | } 410 | } 411 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/aws-saas-genai-rag-workshop/e81b3a60cdb6adaf4d1f9c1fb11e279b279d473d/cdk/lib/tenant-template/services/.DS_Store -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services/aggregate-metrics/invoke_model_tenant_cost.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | import boto3 5 | import time 6 | import os 7 | import json 8 | from datetime import datetime, timedelta 9 | from botocore.exceptions import ClientError 10 | from decimal import * 11 | from aws_lambda_powertools import Tracer, Logger 12 | 13 | tracer = Tracer() 14 | logger = Logger() 15 | 16 | cloudformation = boto3.client('cloudformation') 17 | logs = boto3.client('logs') 18 | athena = boto3.client('athena') 19 | dynamodb = boto3.resource('dynamodb') 20 | # attribution_table = dynamodb.Table("TenantCostAndUsageAttribution") 21 | attribution_table = dynamodb.Table(os.getenv("TENANT_COST_DYNAMODB_TABLE")) 22 | 23 | ATHENA_S3_OUTPUT = os.getenv("ATHENA_S3_OUTPUT") 24 | CUR_DATABASE_NAME = os.getenv("CUR_DATABASE_NAME") 25 | CUR_TABLE_NAME = os.getenv("CUR_TABLE_NAME") 26 | RETRY_COUNT = 100 27 | EMBEDDING_TITAN_INPUT_TOKENS_LABEL="USW2-TitanEmbeddingsG1-Text-input-tokens" 28 | TEXTLITE_INPUT_TOKENS_LABEL="USW2-TitanTextG1-Lite-input-tokens" 29 | TEXTLITE_OUTPUT_TOKENS_LABEL="USW2-TitanTextG1-Lite-output-tokens" 30 | MODEL_INVOCATION_LOG_GROUPNAME= os.getenv("MODEL_INVOCATION_LOG_GROUPNAME") 31 | 32 | class InvokeModelTenantCost(): 33 | 34 | 35 | def __init__(self, start_date_time, end_date_time): 36 | self.start_date_time = start_date_time 37 | self.end_date_time = end_date_time 38 | 39 | def total_service_cost(self): 40 | 41 | # We need to add more filters for day, month, year, resource ids etc. Below query is because we are just using a sample cur file 42 | #Ignoting startTime and endTime filter for now since we have a static/sample cur file 43 | 44 | query = f"SELECT line_item_usage_type, CAST(sum(line_item_blended_cost) AS DECIMAL(10, 6)) AS cost FROM {CUR_DATABASE_NAME}.{CUR_TABLE_NAME} WHERE line_item_product_code='AmazonBedrock' group by 1" 45 | 46 | # Execution 47 | response = athena.start_query_execution( 48 | QueryString=query, 49 | QueryExecutionContext={ 50 | 'Database': CUR_DATABASE_NAME 51 | }, 52 | ResultConfiguration={ 53 | 'OutputLocation': "s3://" + ATHENA_S3_OUTPUT, 54 | } 55 | ) 56 | 57 | # get query execution id 58 | query_execution_id = response['QueryExecutionId'] 59 | logger.info(query_execution_id) 60 | 61 | # get execution status 62 | for i in range(1, 1 + RETRY_COUNT): 63 | 64 | # get query execution 65 | query_status = athena.get_query_execution(QueryExecutionId=query_execution_id) 66 | print (query_status) 67 | query_execution_status = query_status['QueryExecution']['Status']['State'] 68 | 69 | if query_execution_status == 'SUCCEEDED': 70 | print("STATUS:" + query_execution_status) 71 | break 72 | 73 | if query_execution_status == 'FAILED': 74 | raise Exception("STATUS:" + query_execution_status) 75 | 76 | else: 77 | print("STATUS:" + query_execution_status) 78 | time.sleep(i) 79 | else: 80 | athena.stop_query_execution(QueryExecutionId=query_execution_id) 81 | raise Exception('TIME OVER') 82 | 83 | # get query results 84 | result = athena.get_query_results(QueryExecutionId=query_execution_id) 85 | 86 | logger.info (result) 87 | 88 | 89 | 90 | total_service_cost_dict = {} 91 | for row in result['ResultSet']['Rows'][1:]: 92 | line_item = row['Data'][0]['VarCharValue'] 93 | cost = Decimal(row['Data'][1]['VarCharValue']) 94 | # TODO: Lab4 - Get total input and output tokens cost 95 | 96 | 97 | 98 | 99 | logger.info(total_service_cost_dict) 100 | 101 | 102 | # total_service_cost_dict = {"USE1-TitanEmbeddingsG1-Text-input-tokens": 5000, "USE1-TitanTextLite-input-tokens": 4000, "USE1-TitanTextLite-output-tokens": 6000} 103 | return total_service_cost_dict 104 | 105 | def query_metrics(self): 106 | # This dictionary stores the data in format 107 | # {'TenantId': '{"USE1-TitanEmbeddingsG1-Text-input-tokens":0.2, 108 | # "USE1-TitanTextLite-input-tokens": 0.4, "USE1-TitanTextLite-output-tokens": 0.6}'} 109 | tenant_attribution_dict = {} 110 | log_group_names = self.__get_list_of_log_group_names() 111 | 112 | self.__get_tenant_kb_attribution(log_group_names, tenant_attribution_dict) 113 | 114 | self.__get_tenant_converse_attribution(log_group_names, tenant_attribution_dict) 115 | 116 | return tenant_attribution_dict 117 | 118 | 119 | 120 | def calculate_tenant_cost(self, total_service_cost_dict, tenant_attribution_dict): 121 | 122 | for tenant_id, tenant_attribution_percentage in tenant_attribution_dict.items(): 123 | 124 | tenant_attribution_percentage_json = json.loads(tenant_attribution_percentage) 125 | 126 | # TODO: Lab4 - Calculate tenant cost for ingesting & retrieving tenant data to/from Amazon Bedrock Knowledge Base 127 | tenant_kb_input_tokens_cost = 0 128 | 129 | # TODO: Lab4 - Calculate tenant cost for generating final tenant specific response 130 | tenant_input_tokens_cost = 0 131 | tenant_output_tokens_cost = 0 132 | 133 | tenant_service_cost = tenant_kb_input_tokens_cost + tenant_input_tokens_cost + tenant_output_tokens_cost 134 | try: 135 | response = attribution_table.put_item( 136 | Item= 137 | { 138 | "Date": self.start_date_time, 139 | "TenantId#ServiceName": tenant_id+"#"+"AmazonBedrock", 140 | "TenantId": tenant_id, 141 | "TenantKnowledgeBaseInputTokensCost": tenant_kb_input_tokens_cost, 142 | "TenantInputTokensCost": tenant_input_tokens_cost, 143 | "TenantOutputTokensCost": tenant_output_tokens_cost, 144 | "TenantAttributionPercentage": tenant_attribution_percentage, 145 | "TenantServiceCost": tenant_service_cost, 146 | "TotalServiceCost": total_service_cost_dict 147 | } 148 | ) 149 | except ClientError as e: 150 | print(e.response['Error']['Message']) 151 | raise Exception('Error', e) 152 | else: 153 | print("PutItem succeeded:") 154 | 155 | def __is_log_group_exists(self, log_group_name): 156 | logs_paginator = logs.get_paginator('describe_log_groups') 157 | response_iterator = logs_paginator.paginate(logGroupNamePrefix=log_group_name) 158 | for log_groups_list in response_iterator: 159 | if not log_groups_list["logGroups"]: 160 | return False 161 | else: 162 | return True 163 | 164 | def __add_log_group_name(self, log_group_name, log_group_names_list): 165 | if self.__is_log_group_exists(log_group_name): 166 | log_group_names_list.append(log_group_name) 167 | 168 | 169 | def __get_list_of_log_group_names(self): 170 | log_group_names = [] 171 | # Adding bedrock model invocation cloudwatch log group 172 | self.__add_log_group_name(MODEL_INVOCATION_LOG_GROUPNAME, log_group_names) 173 | 174 | # Adding RagService lambda cloudwatch log group 175 | log_group_prefix = '/aws/lambda/' 176 | cloudformation_paginator = cloudformation.get_paginator('list_stack_resources') 177 | response_iterator = cloudformation_paginator.paginate(StackName='saas-genai-workshop-bootstrap-template') 178 | for stack_resources in response_iterator: 179 | for resource in stack_resources['StackResourceSummaries']: 180 | if ("RagService" in resource["LogicalResourceId"] 181 | and resource["ResourceType"] == "AWS::Lambda::Function"): 182 | self.__add_log_group_name(''.join([log_group_prefix,resource["PhysicalResourceId"]]), 183 | log_group_names) 184 | continue 185 | 186 | logger.info(log_group_names) 187 | return log_group_names 188 | 189 | def __query_cloudwatch_logs(self, log_group_names, query_string): 190 | query = logs.start_query(logGroupNames=log_group_names, 191 | startTime=self.start_date_time, 192 | endTime=self.end_date_time, 193 | queryString=query_string) 194 | 195 | query_results = logs.get_query_results(queryId=query["queryId"]) 196 | 197 | while query_results['status']=='Running' or query_results['status']=='Scheduled': 198 | time.sleep(5) 199 | query_results = logs.get_query_results(queryId=query["queryId"]) 200 | 201 | return query_results 202 | 203 | def __get_tenant_kb_attribution(self, log_group_names, tenant_attribution_dict): 204 | 205 | #TODO: Lab4 - Add Amazon CloudWatch logs insights queries to get knowledge base input tokens 206 | knowledgebase_input_tokens_query = "" 207 | 208 | 209 | knowledgebase_input_tokens_resultset = self.__query_cloudwatch_logs(log_group_names, knowledgebase_input_tokens_query) 210 | 211 | 212 | # TODO: Lab4 - Add Amazon CloudWatch logs insights queries to get total knowledge base input tokens 213 | total_knowledgebase_input_tokens_query = "" 214 | 215 | total_knowledgebase_input_tokens_resultset = self.__query_cloudwatch_logs(log_group_names, total_knowledgebase_input_tokens_query) 216 | 217 | # We configure knowledgebase to use Amazon Titan Text Embeddings V2 model to generate embeddings 218 | # When generating embeddings you are only charged for input tokens. 219 | # Here we are calculating the tenant percentage of embedding input tokens 220 | # when interacting with Amazon Titan Text Embeddings V2 model through knowledgebase 221 | if len(total_knowledgebase_input_tokens_resultset['results']) > 0: 222 | total_knowledgebase_input_tokens = Decimal('1') 223 | for row in total_knowledgebase_input_tokens_resultset['results'][0]: 224 | if 'TotalInputTokens' in row['field']: 225 | total_knowledgebase_input_tokens = Decimal(row['value']) 226 | 227 | for row in knowledgebase_input_tokens_resultset['results']: 228 | for field in row: 229 | if 'tenantId' in field['field']: 230 | tenant_id = field['value'] 231 | if 'TotalInputTokens' in field['field']: 232 | input_tokens = Decimal(field['value']) 233 | 234 | # TODO: Lab4 - Calculate the percentage of tenant attribution for knowledge base input tokens 235 | tenant_kb_input_tokens_attribution_percentage = 0 236 | self.__add_or_update_dict(tenant_attribution_dict, tenant_id,EMBEDDING_TITAN_INPUT_TOKENS_LABEL, tenant_kb_input_tokens_attribution_percentage) 237 | 238 | 239 | def __get_tenant_converse_attribution(self, log_group_names, tenant_attribution_dict): 240 | 241 | # TODO: Lab4 - Add Amazon CloudWatch logs insights queries for converse input output tokens 242 | converse_input_output_tokens_query = "" 243 | 244 | converse_input_output_tokens = self.__query_cloudwatch_logs(log_group_names, converse_input_output_tokens_query) 245 | 246 | # TODO: Lab4 - Add Amazon CloudWatch logs insights queries to get total converse input output tokens 247 | total_converse_input_output_tokens_query = "" 248 | 249 | total_converse_input_output_tokens = self.__query_cloudwatch_logs(log_group_names, total_converse_input_output_tokens_query) 250 | 251 | if (len(total_converse_input_output_tokens['results']) > 0): 252 | total_input_tokens = Decimal('1') 253 | total_output_tokens = Decimal('1') 254 | 255 | for row in total_converse_input_output_tokens['results'][0]: 256 | if 'TotalInputTokens' in row['field']: 257 | total_input_tokens = Decimal(row['value']) 258 | if 'TotalOutputTokens' in row['field']: 259 | total_output_tokens = Decimal(row['value']) 260 | 261 | total_input_output_tokens = total_input_tokens + total_output_tokens 262 | 263 | if ( total_input_output_tokens > 0): 264 | 265 | for row in converse_input_output_tokens['results']: 266 | for field in row: 267 | if 'TenantId' in field['field']: 268 | tenant_id = field['value'] 269 | if 'TotalInputTokens' in field['field']: 270 | tenant_input_tokens = Decimal(field['value']) 271 | if 'TotalOutputTokens' in field['field']: 272 | tenant_output_tokens = Decimal(field['value']) 273 | 274 | # TODO: Lab4 - Calculate the percentage of tenant attribution for converse input and output tokens 275 | tenant_attribution_input_tokens_percentage = 0 276 | tenant_attribution_output_tokens_percentage = 0 277 | 278 | self.__add_or_update_dict(tenant_attribution_dict, tenant_id,TEXTLITE_INPUT_TOKENS_LABEL, tenant_attribution_input_tokens_percentage) 279 | self.__add_or_update_dict(tenant_attribution_dict, tenant_id,TEXTLITE_OUTPUT_TOKENS_LABEL, tenant_attribution_output_tokens_percentage) 280 | 281 | 282 | 283 | def __add_or_update_dict(self, tenant_attribution_dict, key, new_attribute_name, new_attribute_value): 284 | if key in tenant_attribution_dict: 285 | # Key exists, so load the JSON string into a Python object 286 | json_obj = json.loads(tenant_attribution_dict[key]) 287 | 288 | # Add the new attribute to the Python object 289 | json_obj[new_attribute_name] = str(new_attribute_value) 290 | 291 | tenant_attribution_dict[key] = json.dumps(json_obj) 292 | else: 293 | # Key does not exist, create a new Python object with the new attribute 294 | new_json_obj = {new_attribute_name: str(new_attribute_value)} 295 | 296 | tenant_attribution_dict[key] = json.dumps(new_json_obj) 297 | 298 | 299 | def __get_tenant_cost(self, key, total_service_cost, tenant_attribution_percentage_json): 300 | tenant_data = tenant_attribution_percentage_json.get(key, 0) 301 | # Bedrock service cost is charged per 1000 tokens 302 | tenant_cost = Decimal(tenant_data) * Decimal(total_service_cost[key])/1000 303 | return tenant_cost -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services/aggregate-metrics/tenant_cost_calculator.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | from datetime import datetime, timedelta 5 | from invoke_model_tenant_cost import InvokeModelTenantCost 6 | from aws_lambda_powertools import Tracer, Logger 7 | 8 | tracer = Tracer() 9 | logger = Logger() 10 | 11 | @tracer.capture_lambda_handler 12 | def calculate_cost_per_tenant(event, context): 13 | 14 | invoke_model_tenant_cost = InvokeModelTenantCost(__get_start_date_time(), 15 | __get_end_date_time()) 16 | 17 | total_service_cost_dict = invoke_model_tenant_cost.total_service_cost() 18 | metrics_dict = invoke_model_tenant_cost.query_metrics() 19 | invoke_model_tenant_cost.calculate_tenant_cost(total_service_cost_dict, metrics_dict) 20 | 21 | 22 | def __get_start_date_time(): 23 | time_zone = datetime.now().astimezone().tzinfo 24 | start_date_time = int(datetime.now(tz=time_zone).date().strftime('%s')) #current day epoch 25 | return start_date_time 26 | 27 | def __get_end_date_time(): 28 | time_zone = datetime.now().astimezone().tzinfo 29 | end_date_time = int((datetime.now(tz=time_zone) + timedelta(days=1)).date().strftime('%s')) #next day epoch 30 | return end_date_time -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services/authorizerService/assume_role_layer.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | import boto3 5 | import hashlib 6 | from aws_lambda_powertools import Logger, Tracer 7 | from collections import namedtuple 8 | 9 | logger = Logger() 10 | 11 | # Define a named tuple for session parameters 12 | SessionParameters = namedtuple( 13 | typename="SessionParameters", 14 | field_names=["aws_access_key_id", "aws_secret_access_key", "aws_session_token"], 15 | ) 16 | 17 | def assume_role(access_role_arn: str, request_tags: list[tuple[str, str]], duration_sec: int = 900) -> SessionParameters: 18 | logger.info(f"Trying to assume role ARN: {access_role_arn} with tags: {request_tags}") 19 | 20 | sts = boto3.client("sts") 21 | 22 | try: 23 | tags_str = "-".join([f"{name}={value}" for name, value in request_tags]) 24 | role_session_name = hashlib.sha256(tags_str.encode()).hexdigest()[:32] 25 | 26 | assume_role_response = sts.assume_role( 27 | RoleArn=access_role_arn, 28 | DurationSeconds=duration_sec, 29 | RoleSessionName=role_session_name, 30 | Tags=[{"Key": name, "Value": value} for name, value in request_tags], 31 | ) 32 | 33 | except Exception as exception: 34 | logger.error(exception) 35 | return None 36 | 37 | logger.info(f"Assumed role ARN: {assume_role_response['AssumedRoleUser']['Arn']}") 38 | 39 | session_parameters = SessionParameters( 40 | aws_access_key_id=assume_role_response["Credentials"]["AccessKeyId"], 41 | aws_secret_access_key=assume_role_response["Credentials"]["SecretAccessKey"], 42 | aws_session_token=assume_role_response["Credentials"]["SessionToken"], 43 | ) 44 | 45 | return session_parameters -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services/authorizerService/authorizer_layer.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | from jose import jwk, jwt 5 | from jose.utils import base64url_decode 6 | import time 7 | import re 8 | import boto3 9 | from aws_lambda_powertools import Logger, Tracer 10 | from collections import namedtuple 11 | 12 | 13 | logger = Logger() 14 | 15 | 16 | def validateJWT(token, app_client_id, keys): 17 | """ 18 | Validate the provided JWT token. 19 | 20 | Args: 21 | token (str): The JWT token to validate. 22 | app_client_id (str): The client ID the token was issued for. 23 | keys (list): The list of public keys to use for verification. 24 | 25 | Returns: 26 | dict: The decoded claims if the token is valid, False otherwise. 27 | """ 28 | try: 29 | # Get the kid from the headers prior to verification 30 | headers = jwt.get_unverified_headers(token) 31 | kid = headers['kid'] 32 | 33 | # Search for the kid in the downloaded public keys 34 | key_index = -1 35 | for i, key in enumerate(keys): 36 | if kid == key['kid']: 37 | key_index = i 38 | break 39 | if key_index == -1: 40 | logger.info('Public key not found in jwks.json') 41 | return False 42 | 43 | # Construct the public key 44 | public_key = jwk.construct(keys[key_index]) 45 | 46 | # Get the last two sections of the token (message and signature) 47 | message, encoded_signature = str(token).rsplit('.', 1) 48 | 49 | # Decode the signature 50 | decoded_signature = base64url_decode(encoded_signature.encode('utf-8')) 51 | 52 | # Verify the signature 53 | if not public_key.verify(message.encode("utf8"), decoded_signature): 54 | logger.info('Signature verification failed') 55 | return False 56 | 57 | logger.info('Signature successfully verified') 58 | 59 | # Get the unverified claims 60 | claims = jwt.get_unverified_claims(token) 61 | 62 | # Verify the token expiration 63 | if time.time() > claims['exp']: 64 | # logger.info('Token is expired') 65 | return False 66 | 67 | # Verify the audience 68 | if claims['aud'] != app_client_id: 69 | logger.info('Token was not issued for this audience') 70 | return False 71 | 72 | logger.info(claims) 73 | return claims 74 | except Exception as e: 75 | logger.error(f"Error validating JWT: {e}") 76 | return False 77 | 78 | class HttpVerb: 79 | GET = "GET" 80 | POST = "POST" 81 | PUT = "PUT" 82 | PATCH = "PATCH" 83 | HEAD = "HEAD" 84 | DELETE = "DELETE" 85 | OPTIONS = "OPTIONS" 86 | ALL = "*" 87 | 88 | class AuthPolicy: 89 | """ 90 | Class for creating and managing authentication policies. 91 | """ 92 | awsAccountId = "" 93 | """The AWS account id the policy will be generated for.""" 94 | principalId = "" 95 | """The principal used for the policy, this should be a unique identifier for the end user.""" 96 | version = "2012-10-17" 97 | """The policy version used for the evaluation.""" 98 | pathRegex = "^[/.a-zA-Z0-9-\*]+$" 99 | """The regular expression used to validate resource paths for the policy""" 100 | 101 | allowMethods = [] 102 | denyMethods = [] 103 | 104 | restApiId = "*" 105 | """The API Gateway API id. By default this is set to '*'""" 106 | region = "*" 107 | """The region where the API is deployed. By default this is set to '*'""" 108 | stage = "*" 109 | """The name of the stage used in the policy. By default this is set to '*'""" 110 | 111 | def __init__(self, principal, awsAccountId): 112 | self.awsAccountId = awsAccountId 113 | self.principalId = principal 114 | self.allowMethods = [] 115 | self.denyMethods = [] 116 | 117 | def _addMethod(self, effect, verb, resource, conditions): 118 | """ 119 | Adds a method to the internal lists of allowed or denied methods. 120 | 121 | Args: 122 | effect (str): "Allow" or "Deny". 123 | verb (str): The HTTP verb (GET, POST, etc.) or "*" for all verbs. 124 | resource (str): The resource path. 125 | conditions (dict): Additional conditions for the policy statement. 126 | 127 | Raises: 128 | NameError: If the HTTP verb or resource path is invalid. 129 | """ 130 | if verb != "*" and not hasattr(HttpVerb, verb): 131 | raise NameError(f"Invalid HTTP verb '{verb}'. Allowed verbs in HttpVerb class") 132 | 133 | resourcePattern = re.compile(self.pathRegex) 134 | if not resourcePattern.match(resource): 135 | raise NameError(f"Invalid resource path: '{resource}'. Path should match '{self.pathRegex}'") 136 | 137 | if resource[:1] == "/": 138 | resource = resource[1:] 139 | 140 | resourceArn = ( 141 | "arn:aws:execute-api:" 142 | + self.region 143 | + ":" 144 | + self.awsAccountId 145 | + ":" 146 | + self.restApiId 147 | + "/" 148 | + self.stage 149 | + "/" 150 | + verb 151 | + "/" 152 | + resource 153 | ) 154 | 155 | if effect.lower() == "allow": 156 | self.allowMethods.append({"resourceArn": resourceArn, "conditions": conditions}) 157 | elif effect.lower() == "deny": 158 | self.denyMethods.append({"resourceArn": resourceArn, "conditions": conditions}) 159 | 160 | # Other methods omitted for brevity 161 | 162 | def create_auth_success_policy( 163 | method_arn: str, 164 | tenant_id: str, 165 | tenant_name: str, 166 | knowledge_base_id: str, 167 | aws_access_key_id: str, 168 | aws_secret_access_key: str, 169 | aws_session_token: str, 170 | api_key: str 171 | ) -> dict: 172 | """ 173 | Creates a success policy for the authorizer to return. 174 | 175 | Args: 176 | method_arn (str): The ARN of the API Gateway method. 177 | tenant_id (str): The tenant ID. 178 | tenant_name (str): The tenant name. 179 | knowledge_base_id (str): The knowledge base ID. 180 | aws_access_key_id (str): The AWS access key ID. 181 | aws_secret_access_key (str): The AWS secret access key. 182 | aws_session_token (str): The AWS session token. 183 | 184 | Returns: 185 | dict: The success policy. 186 | """ 187 | authorization_success_policy = { 188 | "principalId": tenant_id, 189 | "policyDocument": { 190 | "Version": "2012-10-17", 191 | "Statement": [ 192 | { 193 | "Action": "execute-api:Invoke", 194 | "Effect": "Allow", 195 | "Resource": method_arn, 196 | } 197 | ], 198 | }, 199 | "context": { 200 | "knowledge_base_id": knowledge_base_id, 201 | "aws_access_key_id": aws_access_key_id, 202 | "aws_secret_access_key": aws_secret_access_key, 203 | "aws_session_token": aws_session_token, 204 | "tenant_name": tenant_name 205 | }, 206 | "usageIdentifierKey": api_key 207 | } 208 | return authorization_success_policy 209 | 210 | def create_auth_denied_policy(method_arn: str) -> dict: 211 | """ 212 | Creates a deny policy for the authorizer to return. 213 | 214 | Args: 215 | method_arn (str): The ARN of the API Gateway method. 216 | 217 | Returns: 218 | dict: The deny policy. 219 | """ 220 | authorization_deny_policy = { 221 | "principalId": "", 222 | "policyDocument": { 223 | "Version": "2012-10-17", 224 | "Statement": [ 225 | { 226 | "Action": "execute-api:Invoke", 227 | "Effect": "Deny", 228 | "Resource": method_arn, 229 | } 230 | ], 231 | }, 232 | } 233 | return authorization_deny_policy 234 | 235 | 236 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services/authorizerService/tenant_authorizer.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | import urllib.request 5 | import json 6 | import os 7 | import boto3 8 | import time 9 | import re 10 | import authorizer_layer 11 | 12 | from jose import jwk, jwt 13 | from jose.utils import base64url_decode 14 | from aws_lambda_powertools import Logger, Tracer 15 | from collections import namedtuple 16 | from assume_role_layer import assume_role 17 | 18 | 19 | region = os.environ['AWS_REGION'] 20 | userpool_id = os.environ['USER_POOL_ID'] 21 | appclient_id = os.environ['APP_CLIENT_ID'] 22 | assume_role_arn = os.environ['ASSUME_ROLE_ARN'] 23 | control_plane_gw_url = os.environ['CP_API_GW_URL'] 24 | 25 | tenant_token_usage_role_arn = os.environ["TENANT_TOKEN_USAGE_ROLE_ARN"] 26 | tenant_token_usage_table = os.environ["TENANT_TOKEN_USAGE_DYNAMODB_TABLE"] 27 | 28 | logger = Logger() 29 | 30 | def lambda_handler(event, context): 31 | method_arn = event["methodArn"] 32 | 33 | # get JWT token after Bearer from authorization 34 | # Check if authorizationToken is present in the headers 35 | if 'authorizationToken' in event: 36 | token_str = event['authorizationToken'] 37 | elif 'Authorization' in event['headers']: 38 | token_str = event['headers']['Authorization'] 39 | else: 40 | raise Exception('Authorization token is missing') 41 | 42 | token = token_str.split(" ") 43 | if (token[0] != 'Bearer'): 44 | raise Exception('Authorization header should have a format Bearer Token') 45 | jwt_bearer_token = token[1] 46 | 47 | # only to get tenant id to get user pool info 48 | unauthorized_claims = jwt.get_unverified_claims(jwt_bearer_token) 49 | logger.info(unauthorized_claims) 50 | 51 | # get keys for tenant user pool to validate 52 | keys_url = 'https://cognito-idp.{}.amazonaws.com/{}/.well-known/jwks.json'.format(region, userpool_id) 53 | with urllib.request.urlopen(keys_url) as f: 54 | response = f.read() 55 | keys = json.loads(response.decode('utf-8'))['keys'] 56 | 57 | # authenticate against cognito user pool using the key 58 | response = authorizer_layer.validateJWT(jwt_bearer_token, appclient_id, keys) 59 | logger.info(response) 60 | 61 | # get authenticated claims 62 | if (response == False): 63 | logger.error('Unauthorized') 64 | return authorizer_layer.create_auth_denied_policy(method_arn) 65 | else: 66 | tenant_id = response["custom:tenantId"] 67 | response_url = urllib.request.Request(control_plane_gw_url + f'tenant-config?tenantId={tenant_id}') 68 | with urllib.request.urlopen(response_url) as f: 69 | response_data = f.read() 70 | try: 71 | response_text = response_data.decode('utf-8') 72 | response_json = json.loads(response_text) 73 | api_key = response_json['apiKey'] 74 | tenant_name=response_json['tenantName'] 75 | knowledge_base_id = response_json["knowledgeBaseId"] 76 | input_tokens = response_json["inputTokens"] 77 | output_tokens = response_json["outputTokens"] 78 | except UnicodeDecodeError: 79 | print('Unable to decode response data') 80 | except KeyError: 81 | print('API Key not found in response') 82 | 83 | # TODO: Lab3 - Enable tenant token usage 84 | 85 | # assume role 86 | # ABAC: create a temporary session using the assume role arn 87 | # which gives S3 Access to the tenant-specific prefix and specific knowledge base 88 | 89 | request_tags = [("KnowledgeBaseId", knowledge_base_id), ("TenantID", tenant_id)] 90 | session_parameters = assume_role(access_role_arn = assume_role_arn, request_tags = request_tags, duration_sec = 900) 91 | 92 | print("Role Assumed!") 93 | 94 | authorization_success_policy = authorizer_layer.create_auth_success_policy(method_arn, 95 | tenant_id, 96 | tenant_name, 97 | knowledge_base_id, 98 | session_parameters.aws_access_key_id, 99 | session_parameters.aws_secret_access_key, 100 | session_parameters.aws_session_token, 101 | api_key 102 | ) 103 | 104 | logger.debug(authorization_success_policy) 105 | logger.info("Authorization succeeded") 106 | return authorization_success_policy 107 | 108 | 109 | def __is_tenant_token_limit_exceeded(tenant_id, input_tokens, output_tokens): 110 | # check if tenant has enough tokens 111 | try: 112 | input_tokens=int(input_tokens) 113 | output_tokens=int(output_tokens) 114 | table = __get_dynamodb_table(tenant_id) 115 | response = table.get_item( 116 | Key={ 117 | 'TenantId': tenant_id 118 | } 119 | ) 120 | if 'Item' in response: 121 | current_input_tokens = response['Item']['TotalInputTokens'] 122 | current_output_tokens = response['Item']['TotalOutputTokens'] 123 | else: 124 | current_input_tokens = 0 125 | current_output_tokens = 0 126 | 127 | if (current_input_tokens > input_tokens ) or (current_output_tokens > output_tokens ): 128 | logger.error('Tenant token limit exceeded') 129 | return True 130 | 131 | except Exception as e: 132 | logger.error(f"Error validating tenant token usage limit: {e}") 133 | return True 134 | 135 | return False 136 | 137 | def __get_dynamodb_table(tenant_id): 138 | request_tags = [("TenantID", tenant_id)] 139 | session_parameters = assume_role(access_role_arn=tenant_token_usage_role_arn, request_tags = request_tags, duration_sec=900) 140 | dynamodb = boto3.resource('dynamodb', aws_access_key_id=session_parameters.aws_access_key_id, 141 | aws_secret_access_key=session_parameters.aws_secret_access_key, 142 | aws_session_token=session_parameters.aws_session_token) 143 | return dynamodb.Table(tenant_token_usage_table) -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services/layers/metrics_manager.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | import json 5 | from aws_lambda_powertools import Metrics 6 | 7 | metrics = Metrics() 8 | 9 | 10 | def record_metric(event, metric_name, metric_unit, metric_value): 11 | """ Record the metric in Cloudwatch using EMF format 12 | 13 | Args: 14 | event ([type]): [description] 15 | metric_name ([type]): [description] 16 | metric_unit ([type]): [description] 17 | metric_value ([type]): [description] 18 | """ 19 | metrics.add_dimension(name="tenant_id", value=event['requestContext']['authorizer']['principalId']) 20 | metrics.add_metric(name=metric_name, unit=metric_unit, value=metric_value) 21 | metrics_object = metrics.serialize_metric_set() 22 | metrics.clear_metrics() 23 | print(json.dumps(metrics_object)) 24 | 25 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services/layers/requirements.txt: -------------------------------------------------------------------------------- 1 | python-jose[cryptography] 2 | langchain_aws 3 | langchain_core 4 | langchain_community 5 | langchain -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services/ragService/rag_service.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | from aws_lambda_powertools import Logger, Tracer 5 | from aws_lambda_powertools.event_handler import (APIGatewayRestResolver, CORSConfig) 6 | from aws_lambda_powertools.logging import correlation_paths 7 | from aws_lambda_powertools.event_handler.exceptions import ( 8 | InternalServerError, 9 | NotFoundError, 10 | ) 11 | from langchain_aws.chat_models import ChatBedrockConverse 12 | from langchain_core.prompts import ChatPromptTemplate 13 | from langchain_community.retrievers import AmazonKnowledgeBasesRetriever 14 | from langchain.chains import RetrievalQA 15 | from langchain_core.runnables import RunnablePassthrough 16 | import metrics_manager 17 | 18 | import boto3 19 | import json 20 | import os 21 | from botocore.client import Config 22 | 23 | 24 | tracer = Tracer() 25 | logger = Logger() 26 | cors_config = CORSConfig(allow_origin="*", max_age=300) 27 | app = APIGatewayRestResolver(cors=cors_config) 28 | 29 | region_id = os.environ['AWS_REGION'] 30 | MODEL_ID = "amazon.titan-text-lite-v1" 31 | 32 | def retrieveAndGenerate(bedrock_agent_client, bedrock_client, input, knowledge_base_id, event): 33 | vector_retrieval_config={"vectorSearchConfiguration": { 34 | "numberOfResults": 4, 35 | 36 | }} 37 | 38 | retriever = AmazonKnowledgeBasesRetriever( 39 | knowledge_base_id=knowledge_base_id, 40 | retrieval_config=vector_retrieval_config, 41 | client=bedrock_agent_client 42 | ) 43 | 44 | # template = """Answer the question based only on the following context: 45 | # {context} 46 | # Question: {question} 47 | # """ 48 | # chat_prompt_template = ChatPromptTemplate.from_template(template) 49 | 50 | chat_prompt_template = ChatPromptTemplate.from_messages( 51 | [ 52 | 53 | ("human", """Answer the question based only on the following context: 54 | {context} 55 | Question: {question} 56 | Skip a line for each result 57 | """), 58 | 59 | ] 60 | ) 61 | 62 | 63 | llm = ChatBedrockConverse(model=MODEL_ID, temperature=0, client=bedrock_client) 64 | 65 | rag_chain = ( 66 | {"context": retriever , "question": RunnablePassthrough()} 67 | | chat_prompt_template 68 | | llm 69 | ) 70 | 71 | response = rag_chain.invoke(input) 72 | 73 | # Publish metrics 74 | 75 | metrics_manager.record_metric(event, "ModelInvocationInputTokens", "Count", response.usage_metadata['input_tokens']) 76 | metrics_manager.record_metric(event, "ModelInvocationOutputTokens", "Count", response.usage_metadata['output_tokens']) 77 | 78 | 79 | return response.content 80 | 81 | 82 | @logger.inject_lambda_context(correlation_id_path=correlation_paths.API_GATEWAY_REST, log_event=True) 83 | @tracer.capture_lambda_handler 84 | def lambda_handler(event, context): 85 | try: 86 | aws_access_key_id = event['requestContext']['authorizer']['aws_access_key_id'] 87 | aws_secret_access_key = event['requestContext']['authorizer']['aws_secret_access_key'] 88 | aws_session_token = event['requestContext']['authorizer']['aws_session_token'] 89 | knowledge_base_id = event['requestContext']['authorizer']['knowledge_base_id'] 90 | tenant_name = event['requestContext']['authorizer']['tenant_name'] 91 | logger.info(f"input tenant name: {tenant_name} and its knowledge_base_id: {knowledge_base_id}") 92 | # TODO: Lab2 - uncomment below and hardcode an knowledge base id 93 | # knowledge_base_id = "" 94 | # logger.info(f"hard coded knowledge base id: {knowledge_base_id}") 95 | 96 | 97 | if 'body' not in event: 98 | raise ValueError('No query provided') 99 | 100 | # Extract the query from the event 101 | query = event['body'] 102 | 103 | # Log the body content 104 | logger.debug("Received query:", query) 105 | 106 | session = boto3.Session( 107 | aws_access_key_id = aws_access_key_id, 108 | aws_secret_access_key = aws_secret_access_key, 109 | aws_session_token = aws_session_token 110 | ) 111 | 112 | # Initialize the Bedrock client 113 | bedrock_config = Config(connect_timeout=120, read_timeout=120, retries={'max_attempts': 0}) 114 | bedrock_agent_client = session.client('bedrock-agent-runtime', config=bedrock_config) 115 | bedrock_client = session.client('bedrock-runtime', config=bedrock_config) 116 | 117 | response = retrieveAndGenerate(bedrock_agent_client, bedrock_client, query, knowledge_base_id, event) 118 | 119 | logger.info(f"Used the knowledge_base_id: {knowledge_base_id} to generate this response: {response}") 120 | 121 | # Return the results 122 | return { 123 | 'statusCode': 200, 124 | 'body': json.dumps(response) 125 | } 126 | 127 | except Exception as e: 128 | logger.exception("An error occurred during execution") 129 | return { 130 | 'statusCode': 500, 131 | 'body': json.dumps({ 132 | 'error': str(e), 133 | 'message': 'An error occurred during execution' 134 | }) 135 | } -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services/s3Uploader/s3uploader.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | import json 5 | import boto3 6 | import os 7 | 8 | from datetime import datetime 9 | from aws_lambda_powertools import Logger, Tracer 10 | from aws_lambda_powertools.event_handler import (APIGatewayRestResolver, CORSConfig) 11 | from aws_lambda_powertools.logging import correlation_paths 12 | from aws_lambda_powertools.event_handler.exceptions import ( 13 | InternalServerError, 14 | NotFoundError, 15 | ) 16 | 17 | s3_bucket = os.environ['S3_BUCKET_NAME'] 18 | 19 | tracer = Tracer() 20 | logger = Logger() 21 | cors_config = CORSConfig(allow_origin="*", max_age=300) 22 | app = APIGatewayRestResolver(cors=cors_config) 23 | 24 | @logger.inject_lambda_context(correlation_id_path=correlation_paths.API_GATEWAY_REST, log_event=True) 25 | @tracer.capture_lambda_handler 26 | def lambda_handler(event, context): 27 | logger.debug(event) 28 | 29 | # Retrieve tenantId from event 30 | tenantId = event['requestContext']['authorizer']['principalId'] 31 | aws_access_key_id = event['requestContext']['authorizer']['aws_access_key_id'] 32 | aws_secret_access_key = event['requestContext']['authorizer']['aws_secret_access_key'] 33 | aws_session_token = event['requestContext']['authorizer']['aws_session_token'] 34 | 35 | # Assume S3 prefix is set to tenant_id 36 | s3_prefix = tenantId 37 | 38 | # Extract the notes from the event 39 | if 'body' in event: 40 | notes = event['body'] 41 | # Log the body content 42 | logger.debug("Received notes:", notes) 43 | else: 44 | notes = "No notes found" 45 | 46 | session = boto3.Session( 47 | aws_access_key_id = aws_access_key_id, 48 | aws_secret_access_key = aws_secret_access_key, 49 | aws_session_token = aws_session_token 50 | ) 51 | 52 | # Initialize S3 client 53 | s3 = session.client('s3') 54 | 55 | # Generate a unique file name with current time to avoid duplication 56 | current_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") 57 | file_name = f"meeting_notes_{current_time}.txt" 58 | 59 | # Upload the meeting notes to the S3 prefix 60 | s3_key = f"{s3_prefix}/{file_name}" 61 | 62 | try: 63 | s3.put_object(Bucket=s3_bucket, Key=s3_key, Body=notes) 64 | logger.info(f"Meeting notes uploaded to S3://{s3_prefix}/{file_name}") 65 | 66 | # Return a success response 67 | return { 68 | 'statusCode': 200, 69 | 'body': json.dumps({ 70 | 'message': 'Received your notes', 71 | }) 72 | } 73 | 74 | except s3.exceptions.NoSuchBucket as e: 75 | logger.error(f"Bucket not found: {e}") 76 | return { 77 | 'statusCode': 404, 78 | 'body': json.dumps({ 79 | 'message': 'Bucket not found' 80 | }) 81 | } 82 | except s3.exceptions.ClientError as e: 83 | logger.error(f"Client error when uploading meeting notes to S3: {e}") 84 | return { 85 | 'statusCode': 400, 86 | 'body': json.dumps({ 87 | 'message': 'Client error when uploading to S3', 88 | 'error': str(e) 89 | }) 90 | } 91 | except Exception as e: 92 | logger.error(f"Error uploading meeting notes to S3: {e}") 93 | return { 94 | 'statusCode': 500, 95 | 'body': json.dumps({ 96 | 'message': 'Internal server error', 97 | 'error': str(e) 98 | }) 99 | } 100 | 101 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services/tenant-token-usage/tenant_token_usage_calculator.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | from datetime import datetime, timedelta 5 | from aws_lambda_powertools import Tracer, Logger 6 | import boto3 7 | import os 8 | import time 9 | from decimal import Decimal 10 | tracer = Tracer() 11 | logger = Logger() 12 | 13 | 14 | RAG_SERVICE_CLOUDWATCH_LOGS= "/aws/lambda/ragService" 15 | logs = boto3.client('logs') 16 | athena = boto3.client('athena') 17 | dynamodb = boto3.resource('dynamodb') 18 | tenant_token_usage_table = dynamodb.Table(os.getenv("TENANT_TOKEN_USAGE_DYNAMODB_TABLE")) 19 | 20 | 21 | @tracer.capture_lambda_handler 22 | def calculate_daily_tenant_token_usage(event, context): 23 | log_group_names = [RAG_SERVICE_CLOUDWATCH_LOGS] 24 | 25 | tenant_token_usage_query = "filter @message like /ModelInvocationInputTokens|ModelInvocationOutputTokens/ \ 26 | | fields tenant_id as TenantId, ModelInvocationInputTokens.0 as ModelInvocationInputTokens, ModelInvocationOutputTokens.0 as ModelInvocationOutputTokens \ 27 | | stats sum(ModelInvocationInputTokens) as TotalInputTokens, sum(ModelInvocationOutputTokens) as TotalOutputTokens by TenantId, dateceil(@timestamp, 1d) as timestamp" 28 | 29 | tenant_token_usage_resultset = __query_cloudwatch_logs(log_group_names, tenant_token_usage_query) 30 | 31 | logger.info(f'Returned tenant token usage result set of size: {len(tenant_token_usage_resultset['results'])}') 32 | 33 | if len(tenant_token_usage_resultset['results']) > 0: 34 | for row in tenant_token_usage_resultset['results']: 35 | for field in row: 36 | if 'TenantId' in field['field']: 37 | tenant_id = field['value'] 38 | if 'TotalInputTokens' in field['field']: 39 | total_input_tokens = Decimal(field['value']) 40 | if 'TotalOutputTokens' in field['field']: 41 | total_output_tokens = Decimal(field['value']) 42 | 43 | 44 | tenant_token_usage_table.put_item( 45 | Item={ 46 | 'TenantId': tenant_id, 47 | 'TotalInputTokens': total_input_tokens, 48 | 'TotalOutputTokens': total_output_tokens, 49 | 'StartDate': __get_start_date_time(), 50 | 'EndDate': __get_end_date_time() 51 | } 52 | ) 53 | 54 | 55 | 56 | 57 | def __get_start_date_time(): 58 | time_zone = datetime.now().astimezone().tzinfo 59 | start_date_time = int(datetime.now(tz=time_zone).date().strftime('%s')) #current day epoch 60 | return start_date_time 61 | 62 | def __get_end_date_time(): 63 | time_zone = datetime.now().astimezone().tzinfo 64 | end_date_time = int((datetime.now(tz=time_zone) + timedelta(days=1)).date().strftime('%s')) #next day epoch 65 | return end_date_time 66 | 67 | 68 | def __query_cloudwatch_logs(log_group_names, query_string): 69 | query = logs.start_query(logGroupNames=log_group_names, 70 | startTime=__get_start_date_time(), 71 | endTime=__get_end_date_time(), 72 | queryString=query_string) 73 | 74 | query_results = logs.get_query_results(queryId=query["queryId"]) 75 | 76 | while query_results['status']=='Running' or query_results['status']=='Scheduled': 77 | time.sleep(5) 78 | query_results = logs.get_query_results(queryId=query["queryId"]) 79 | 80 | return query_results 81 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services/triggerDataIngestionService/assume_role_layer.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | import boto3 5 | import hashlib 6 | from aws_lambda_powertools import Logger, Tracer 7 | from collections import namedtuple 8 | 9 | logger = Logger() 10 | 11 | # Define a named tuple for session parameters 12 | SessionParameters = namedtuple( 13 | typename="SessionParameters", 14 | field_names=["aws_access_key_id", "aws_secret_access_key", "aws_session_token"], 15 | ) 16 | 17 | def assume_role(access_role_arn: str, request_tags: list[tuple[str, str]], duration_sec: int = 900) -> SessionParameters: 18 | logger.info(f"Trying to assume role ARN: {access_role_arn} with tags: {request_tags}") 19 | 20 | sts = boto3.client("sts") 21 | 22 | try: 23 | tags_str = "-".join([f"{name}={value}" for name, value in request_tags]) 24 | role_session_name = hashlib.sha256(tags_str.encode()).hexdigest()[:32] 25 | 26 | assume_role_response = sts.assume_role( 27 | RoleArn=access_role_arn, 28 | DurationSeconds=duration_sec, 29 | RoleSessionName=role_session_name, 30 | Tags=[{"Key": name, "Value": value} for name, value in request_tags], 31 | ) 32 | 33 | except Exception as exception: 34 | logger.error(exception) 35 | return None 36 | 37 | logger.info(f"Assumed role ARN: {assume_role_response['AssumedRoleUser']['Arn']}") 38 | 39 | session_parameters = SessionParameters( 40 | aws_access_key_id=assume_role_response["Credentials"]["AccessKeyId"], 41 | aws_secret_access_key=assume_role_response["Credentials"]["SecretAccessKey"], 42 | aws_session_token=assume_role_response["Credentials"]["SessionToken"], 43 | ) 44 | 45 | return session_parameters -------------------------------------------------------------------------------- /cdk/lib/tenant-template/services/triggerDataIngestionService/trigger_data_ingestion.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | import json 5 | import boto3 6 | import os 7 | 8 | from datetime import datetime 9 | from assume_role_layer import assume_role 10 | from aws_lambda_powertools import Logger, Tracer 11 | 12 | from aws_lambda_powertools.event_handler import (APIGatewayRestResolver, 13 | CORSConfig) 14 | from aws_lambda_powertools.logging import correlation_paths 15 | from aws_lambda_powertools.event_handler.exceptions import ( 16 | InternalServerError, 17 | NotFoundError, 18 | ) 19 | 20 | tracer = Tracer() 21 | logger = Logger() 22 | cors_config = CORSConfig(allow_origin="*", max_age=300) 23 | app = APIGatewayRestResolver(cors=cors_config) 24 | assume_role_arn = os.environ['ASSUME_ROLE_ARN'] 25 | 26 | @logger.inject_lambda_context(correlation_id_path=correlation_paths.API_GATEWAY_REST, log_event=True) 27 | @tracer.capture_lambda_handler 28 | 29 | def lambda_handler(event, context): 30 | try: 31 | # Log the entire event object for debugging purposes 32 | logger.debug(f"Received event: {event}") 33 | 34 | # Extract kb_id,bucket, and key from the event 35 | 36 | knowledge_base_id = event['kb_id'] 37 | datasource_id = event['datasource_id'] 38 | bucket = event['bucket'] 39 | key = event['key'] 40 | tenant_id = key.split('/')[0] 41 | 42 | # ABAC: create a temporary session using the assume role arn which gives S3 Access to the tenant-specific prefix 43 | 44 | logger.info(tenant_id) 45 | 46 | request_tags = [("KnowledgeBaseId", knowledge_base_id)] 47 | session_parameters = assume_role(access_role_arn=assume_role_arn, request_tags = request_tags, duration_sec = 900) 48 | 49 | session = boto3.Session( 50 | aws_access_key_id=session_parameters.aws_access_key_id, 51 | aws_secret_access_key=session_parameters.aws_secret_access_key, 52 | aws_session_token=session_parameters.aws_session_token 53 | ) 54 | 55 | client = session.client('bedrock-agent') 56 | 57 | response = client.start_ingestion_job( 58 | dataSourceId=datasource_id, 59 | description='data source updated', 60 | knowledgeBaseId=knowledge_base_id 61 | ) 62 | logger.info(response) 63 | 64 | 65 | except KeyError as e: 66 | # Log the error and return a meaningful response 67 | logger.error(f"KeyError: {e}. Event: {event}") 68 | return { 69 | 'statusCode': 400, 70 | 'body': json.dumps('Bad Request: Missing required key') 71 | } 72 | except Exception as e: 73 | # Log any other exceptions and return a meaningful response 74 | logger.error(f"Unhandled exception: {e}. Event: {event}") 75 | return { 76 | 'statusCode': 500, 77 | 'body': json.dumps('Internal Server Error') 78 | } 79 | 80 | # Rest of your code 81 | return { 82 | 'statusCode': 200, 83 | 'body': json.dumps('Success') 84 | } -------------------------------------------------------------------------------- /cdk/lib/tenant-template/tenant-provisioning/requirements.txt: -------------------------------------------------------------------------------- 1 | opensearch-py 2 | boto3>=1.34.91 3 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/tenant-provisioning/tenant_provisioning_service.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | import boto3 5 | import json 6 | import argparse 7 | import time 8 | from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth 9 | import os 10 | import logging 11 | import uuid 12 | import sys 13 | 14 | 15 | # HOST=os.environ['OPENSEARCH_SERVERLESS_ENDPOINT'] 16 | REGION=os.environ['AWS_REGION'] 17 | S3_BUCKET=os.environ['S3_BUCKET'] 18 | TRIGGER_PIPELINE_INGESTION_LAMBDA_ARN=os.environ['TRIGGER_PIPELINE_INGESTION_LAMBDA_ARN'] 19 | OPENSEARCH_SERVERLESS_COLLECTION_ARN=os.environ['OPENSEARCH_SERVERLESS_COLLECTION_ARN'] 20 | TENANT_API_KEY=os.environ['TENANT_API_KEY'] 21 | EMBEDDING_MODEL_ARN = f'arn:aws:bedrock:{REGION}::foundation-model/amazon.titan-embed-text-v1' 22 | TENANT_KB_METADATA_FIELD = 'tenant-knowledge-base-metadata' 23 | TENANT_KB_TEXT_FIELD = 'tenant-knowledge-base-text' 24 | TENANT_KB_VECTOR_FIELD = 'tenant-knowledge-base-vector' 25 | 26 | aoss_client = boto3.client('opensearchserverless') 27 | s3 = boto3.client('s3') 28 | bedrock_agent_client = boto3.client('bedrock-agent') 29 | iam_client = boto3.client('iam') 30 | eventbridge = boto3.client('events') 31 | lambda_client = boto3.client('lambda') 32 | 33 | def provision_tenant_resources(tenant_id): 34 | 35 | kb_collection = __get_opensearch_serverless_collection_details() 36 | kb_collection_name = kb_collection['name'] 37 | kb_collection_endpoint = kb_collection['collectionEndpoint'] 38 | kb_collection_endpoint_domain=kb_collection_endpoint.split("//")[-1] 39 | rule_name=f's3rule-{tenant_id}' 40 | 41 | try: 42 | # TODO: Lab1 - Add provision tenant resources 43 | 44 | __api_gw_add_api_key(tenant_id) 45 | return 0 46 | except Exception as e: 47 | logging.error('Error occured while provisioning tenant resources', e) 48 | return 1 49 | 50 | # Function adding api key to existing api gateway usage plan 51 | def __api_gw_add_api_key(tenant_id): 52 | try: 53 | api_key_value = TENANT_API_KEY 54 | usage_plan_id = os.environ['API_GATEWAY_USAGE_PLAN_ID'] 55 | apigw_client = boto3.client('apigateway') 56 | response = apigw_client.create_api_key( 57 | name=tenant_id, 58 | description='Tenant API Key', 59 | enabled=True, 60 | value=api_key_value 61 | ) 62 | api_key = response['id'] 63 | apigw_client.create_usage_plan_key( 64 | usagePlanId=usage_plan_id, 65 | keyId=api_key, 66 | keyType='API_KEY' 67 | ) 68 | logging.info(f'API key {api_key} added to usage plan {usage_plan_id}') 69 | return 0 70 | except Exception as e: 71 | logging.error('Error occured while adding api key to api gateway usage plan', e) 72 | return 1 73 | 74 | def __get_opensearch_serverless_collection_details(): 75 | try: 76 | kb_collection_id = OPENSEARCH_SERVERLESS_COLLECTION_ARN.split('/')[-1] 77 | kb_collection = aoss_client.batch_get_collection( 78 | ids=[kb_collection_id] 79 | ) 80 | 81 | logging.info(f'OpenSearch serverless collection details: {kb_collection}') 82 | return kb_collection['collectionDetails'][0] 83 | except Exception as e: 84 | logging.error('Error occured while getting OpenSearch serverless collection details', e) 85 | raise Exception('Error occured while getting OpenSearch serverless collection details') from e 86 | 87 | def __create_tenant_knowledge_base(tenant_id, kb_collection_name, rule_name): 88 | try: 89 | tenant_kb_role_arn = __create_tenant_kb_role(tenant_id) 90 | 91 | storage_configuration = { 92 | 'opensearchServerlessConfiguration': { 93 | 'collectionArn': OPENSEARCH_SERVERLESS_COLLECTION_ARN, 94 | 'fieldMapping': { 95 | 'metadataField': TENANT_KB_METADATA_FIELD, 96 | 'textField': TENANT_KB_TEXT_FIELD, 97 | 'vectorField': TENANT_KB_VECTOR_FIELD 98 | }, 99 | 'vectorIndexName': tenant_id 100 | }, 101 | 'type': 'OPENSEARCH_SERVERLESS' 102 | } 103 | 104 | 105 | __add_data_access_policy(tenant_id, tenant_kb_role_arn, kb_collection_name) 106 | 107 | # Wait for the IAM role to be created 108 | logging.info(f'Waiting for IAM role "bedrock-kb-role-{tenant_id}" to be created...') 109 | time.sleep(10) 110 | 111 | # Retries to handle delay in indexes or Data Access Policies to be available. Indexes and Data Access Policy could take few seconds to be available to Bedrock KB. 112 | num_retries = 0 113 | max_retries = 10 114 | while num_retries < max_retries: 115 | try: 116 | response = bedrock_agent_client.create_knowledge_base( 117 | name=tenant_id, 118 | description=f'Knowledge base for tenant {tenant_id}', 119 | roleArn=tenant_kb_role_arn, 120 | knowledgeBaseConfiguration={ 121 | 'type': 'VECTOR', 122 | 'vectorKnowledgeBaseConfiguration': { 123 | 'embeddingModelArn': EMBEDDING_MODEL_ARN 124 | }, 125 | }, 126 | storageConfiguration=storage_configuration 127 | ); 128 | logging.info(f'Tenant knowledge base created: {response}') 129 | except bedrock_agent_client.exceptions.ValidationException as e: 130 | error_message = e.response['Error']['Message'] 131 | logging.error(f'{error_message}. Retrying in 5 seconds') 132 | time.sleep(5) 133 | num_retries += 1 134 | except bedrock_agent_client.exceptions.ConflictException: 135 | logging.info(f"Knowledge base '{tenant_id}' already exists, skipping creation.") 136 | break 137 | except Exception as e: 138 | logging.error('Error occurred while creating tenant knowledge base', e) 139 | raise Exception('Error occurred while creating tenant knowledge base') from e 140 | else: 141 | logging.error('Maximum number of retries reached, giving up.') 142 | raise Exception('Error occurred while creating tenant knowledge base: Maximum number of retries reached') 143 | knowledge_base_id = response['knowledgeBase']['knowledgeBaseId'] 144 | logging.info(knowledge_base_id) 145 | datasource_id = __create_tenant_data_source(tenant_id, knowledge_base_id) 146 | __create_eventbridge_tenant_rule_target(tenant_id, knowledge_base_id, rule_name, datasource_id) 147 | 148 | except Exception as e: 149 | logging.error('Error occured while creating tenant knowledge base', e) 150 | raise Exception('Error occured while creating tenant knowledge base') from e 151 | 152 | # Create Data Source in S3 per Tenant 153 | def __create_tenant_data_source(tenant_id, knowledge_base_id): 154 | s3_bucket_arn=f'arn:aws:s3:::{S3_BUCKET}' 155 | response = bedrock_agent_client.create_data_source( 156 | name=tenant_id, 157 | description=f'Data source for tenant {tenant_id}', 158 | dataSourceConfiguration={ 159 | 's3Configuration':{ 160 | 'bucketArn':s3_bucket_arn, 161 | 'inclusionPrefixes': [f'{tenant_id}/'] 162 | }, 163 | 'type': 'S3' 164 | }, 165 | knowledgeBaseId=knowledge_base_id 166 | ) 167 | 168 | return response['dataSource']['dataSourceId'] 169 | 170 | 171 | def __create_tenant_kb_role(tenant_id): 172 | try: 173 | try: 174 | response = iam_client.create_role( 175 | RoleName=f'bedrock-kb-role-{tenant_id}', 176 | AssumeRolePolicyDocument=json.dumps(__get_kb_trust_policy())) 177 | except Exception as e: 178 | if e.response['Error']['Code'] == 'EntityAlreadyExists': 179 | logging.info (f"IAM role 'bedrock-kb-role-{tenant_id}' already exists, skipping creation.") 180 | response = iam_client.get_role( 181 | RoleName=f'bedrock-kb-role-{tenant_id}') 182 | 183 | iam_client.put_role_policy( 184 | RoleName=f'bedrock-kb-role-{tenant_id}', 185 | PolicyName=f'bedrock-kb-policy-{tenant_id}', 186 | PolicyDocument=json.dumps(__get_kb_policy(tenant_id)) 187 | ) 188 | logging.info(f"Tenant knowledge base role created: {response['Role']['Arn']}") 189 | return response['Role']['Arn'] 190 | 191 | except Exception as e: 192 | logging.error('Error occured while creating tenant knowledge base role', e) 193 | raise Exception('Error occured while creating tenant knowledge base role') from e 194 | 195 | # Function creating a Data Access Policy per tenant 196 | def __add_data_access_policy(tenant_id, tenant_kb_role_arn, kb_collection_name): 197 | 198 | # Trimming tenant id to accomodate the polic name 32 characters limit 199 | # Shortening tenant id to 25 characters to fit the policy name 200 | trimmed_tenant_id = tenant_id[:25] 201 | 202 | try: 203 | 204 | response=aoss_client.create_access_policy( 205 | name=f'policy-{trimmed_tenant_id}', 206 | description=f'Data Access Policy for tenant {tenant_id}', 207 | policy=json.dumps(__generate_data_access_policy(tenant_id, tenant_kb_role_arn, kb_collection_name)), 208 | type='data') 209 | 210 | logging.info(f'Tenant data access policy created: {response}') 211 | 212 | except Exception as e: 213 | logging.error('Error occured while adding data access policy', e) 214 | raise Exception('Error occured while adding data access policy') from e 215 | 216 | def __create_s3_tenant_prefix(tenant_id, rule_name): 217 | try: 218 | prefix = ''.join([tenant_id, '/']) 219 | s3.put_object(Bucket=S3_BUCKET, Key=prefix) 220 | rule_arn=__create_eventbridge_tenant_rule(prefix, tenant_id, rule_name) 221 | __create_trigger_lambda_eventbridge_permissions(rule_arn) 222 | logging.info(f'S3 tenant prefix created for tenant {tenant_id}') 223 | return rule_name 224 | 225 | except Exception as e: 226 | logging.error('Error occured while creating S3 tenant prefix', e) 227 | raise Exception('Error occured while creating S3 tenant prefix') from e 228 | 229 | def __create_eventbridge_tenant_rule(prefix, tenant_id, rule_name): 230 | try: 231 | event_pattern = { 232 | "detail": { 233 | "bucket": { 234 | "name": [S3_BUCKET] 235 | }, 236 | "object": { 237 | "key": [{ 238 | "prefix": prefix 239 | }] 240 | } 241 | }, 242 | "detail-type": ["Object Created"], 243 | "source": ["aws.s3"] 244 | } 245 | 246 | rule = eventbridge.put_rule( 247 | Name=rule_name, 248 | EventPattern=json.dumps(event_pattern), 249 | State='ENABLED' 250 | ) 251 | 252 | logging.info(f'Eventbridge Rule Created for tenant {tenant_id}') 253 | 254 | return rule['RuleArn'] 255 | 256 | except Exception as e: 257 | logging.error('Error occured while creating eventbridge rule', e) 258 | raise Exception('Error occured while creating eventbridge rule', e) 259 | 260 | def __create_trigger_lambda_eventbridge_permissions(rule_arn): 261 | try: 262 | lambda_client.add_permission( 263 | FunctionName=TRIGGER_PIPELINE_INGESTION_LAMBDA_ARN, 264 | StatementId=f'bedrock-pipeline-ingestion-{uuid.uuid4()}', 265 | Action='lambda:InvokeFunction', 266 | Principal='events.amazonaws.com', 267 | SourceArn=rule_arn 268 | ) 269 | logging.info(f'Trigger Lambda EventBridge permissions created') 270 | 271 | except Exception as e: 272 | logging.error('Error occured while creating trigger lambda eventbridge permissions', e) 273 | raise Exception('Error occured while creating trigger lambda eventbridge permissions', e) 274 | 275 | def __create_eventbridge_tenant_rule_target(tenant_id, kb_id, rule_name, datasource_id): 276 | try: 277 | # input_template=f'{{"kb_id": "{kb_id}", "datasource_id": "{datasource_id}", "bucket": , "key": }}' 278 | input_template = { 279 | "kb_id": kb_id, 280 | "datasource_id": datasource_id, 281 | "bucket": "", 282 | "key": "" 283 | } 284 | input_transformer = { 285 | 'InputPathsMap': { 286 | "object-key": "$.detail.object.key", 287 | "bucket": "$.detail.bucket.name" 288 | }, 289 | "InputTemplate": json.dumps(input_template) 290 | } 291 | 292 | eventbridge.put_targets( 293 | Rule=rule_name, 294 | Targets=[ 295 | { 296 | 'Id': tenant_id, 297 | 'Arn': TRIGGER_PIPELINE_INGESTION_LAMBDA_ARN, 298 | 'InputTransformer': input_transformer, 299 | 'RetryPolicy': { 300 | 'MaximumRetryAttempts': 2, 301 | 'MaximumEventAgeInSeconds': 3600 302 | } 303 | } 304 | ] 305 | ) 306 | logging.info(f'Eventbridge rule target created for tenant {tenant_id}') 307 | 308 | except Exception as e: 309 | logging.error('Error occured while creating eventbridge rule', e) 310 | raise Exception('Error occured while creating eventbridge rule') from e 311 | 312 | def __create_opensearch_serverless_tenant_index(tenantId, kb_collection_endpoint): 313 | try: 314 | # Get AWS credentials 315 | credentials = boto3.Session().get_credentials() 316 | 317 | # Create the AWS Signature Version 4 signer for OpenSearch Serverless 318 | auth = AWSV4SignerAuth(credentials, REGION, 'aoss') 319 | 320 | # Create the OpenSearch client 321 | client = OpenSearch( 322 | hosts=[{'host': kb_collection_endpoint, 'port': 443}], 323 | http_auth=auth, 324 | use_ssl=True, 325 | verify_certs=True, 326 | connection_class=RequestsHttpConnection, 327 | pool_maxsize=20 328 | ) 329 | 330 | index_body = { 331 | "settings": { 332 | "index.knn": True, 333 | "number_of_shards": 1, 334 | "knn.algo_param.ef_search": 512, 335 | "number_of_replicas": 0, 336 | }, 337 | "mappings": { 338 | "properties": {} 339 | } 340 | } 341 | 342 | index_body["mappings"]["properties"][TENANT_KB_VECTOR_FIELD] = { 343 | "type": "knn_vector", 344 | "dimension": 1536, 345 | "method": { 346 | "name": "hnsw", 347 | "engine": "faiss" 348 | }, 349 | } 350 | 351 | index_body["mappings"]["properties"][TENANT_KB_TEXT_FIELD] = { 352 | "type": "text" 353 | } 354 | 355 | index_body["mappings"]["properties"][TENANT_KB_METADATA_FIELD] = { 356 | "type": "text" 357 | } 358 | 359 | # Create the index 360 | try: 361 | response = client.indices.create(index=tenantId, body=index_body) 362 | logging.info(f'Tenant open search serverless index created: {response}') 363 | except Exception as e: 364 | if 'resource_already_exists_exception' in str(e).lower(): 365 | logging.info(f'Tenant open search serverless index {tenantId} already exists, skipping creation.') 366 | else: 367 | logging.error('Error occurred while creating opensearch serverless tenant index', e) 368 | raise Exception('Error occurred while creating opensearch serverless tenant index') from e 369 | except Exception as e: 370 | logging.error('Error occured while creating opensearch serverless tenant prefix', e) 371 | raise Exception('Error occured while creating opensearch serverless tenant prefix') from e 372 | 373 | 374 | 375 | def __get_kb_trust_policy(): 376 | return { 377 | "Version": "2012-10-17", 378 | "Statement": [{ 379 | "Effect": "Allow", 380 | "Principal": { 381 | "Service": "bedrock.amazonaws.com" 382 | }, 383 | "Action": "sts:AssumeRole" 384 | }] 385 | } 386 | 387 | def __get_kb_policy(tenant_id): 388 | return { 389 | "Version": "2012-10-17", 390 | "Statement": [ 391 | { #FM model policy 392 | "Sid": "AmazonBedrockAgentBedrockFoundationModelPolicy", 393 | "Effect": "Allow", 394 | "Action": "bedrock:InvokeModel", 395 | "Resource": [ 396 | EMBEDDING_MODEL_ARN 397 | ] 398 | }, 399 | { #AOSS policy 400 | "Effect": "Allow", 401 | "Action": "aoss:APIAccessAll", 402 | "Resource": [OPENSEARCH_SERVERLESS_COLLECTION_ARN] 403 | }, 404 | { # S3 policy 405 | "Sid": "AllowKBAccessDocuments", 406 | "Effect": "Allow", 407 | "Action": [ 408 | "s3:GetObject" 409 | ], 410 | "Resource": [f"arn:aws:s3:::{S3_BUCKET}/{tenant_id}/*"] 411 | }, 412 | { 413 | "Sid": "AllowKBAccessBucket", 414 | "Effect": "Allow", 415 | "Action": [ 416 | "s3:ListBucket" 417 | ], 418 | "Resource": [ 419 | f"arn:aws:s3:::{S3_BUCKET}" 420 | ], 421 | "Condition": { 422 | "StringLike": { 423 | "s3:prefix": [ 424 | f"{tenant_id}/*" 425 | ] 426 | } 427 | } 428 | } 429 | ] 430 | } 431 | 432 | # Generating a Data Access Policy per tenant 433 | def __generate_data_access_policy(tenant_id, tenant_kb_role_arn, kb_collection_name): 434 | return [ 435 | { 436 | "Rules": [ 437 | { 438 | "Resource": [ 439 | f"index/{kb_collection_name}/{tenant_id}" 440 | ], 441 | "Permission": [ 442 | "aoss:CreateIndex", 443 | "aoss:DeleteIndex", 444 | "aoss:UpdateIndex", 445 | "aoss:DescribeIndex", 446 | "aoss:ReadDocument", 447 | "aoss:WriteDocument" 448 | ], 449 | "ResourceType": "index" 450 | } 451 | ], 452 | "Principal": [tenant_kb_role_arn], 453 | } 454 | ] 455 | 456 | 457 | 458 | if __name__ == '__main__': 459 | parser = argparse.ArgumentParser( 460 | description='Provisioning tenant resources') 461 | parser.add_argument('--tenantid', type=str, help='tenantid', required=True) 462 | args = parser.parse_args() 463 | # tenantId=args.tenantid.strip('"') 464 | 465 | # provision_tenant_resources(**vars(args)) 466 | status = provision_tenant_resources(args.tenantid) 467 | sys.exit(status) 468 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/tenant-token-usage.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import { Construct } from "constructs"; 5 | import * as cdk from "aws-cdk-lib"; 6 | import { PythonFunction } from "@aws-cdk/aws-lambda-python-alpha"; 7 | import * as iam from "aws-cdk-lib/aws-iam"; 8 | import * as lambda from "aws-cdk-lib/aws-lambda"; 9 | import * as dynamodb from "aws-cdk-lib/aws-dynamodb"; 10 | import * as events from "aws-cdk-lib/aws-events"; 11 | import * as targets from "aws-cdk-lib/aws-events-targets"; 12 | 13 | import * as path from "path"; 14 | 15 | export interface TenantTokenUsageProps { 16 | readonly lambdaPowerToolsLayer: lambda.ILayerVersion; 17 | } 18 | 19 | export class TenantTokenUsage extends Construct { 20 | readonly tenantTokenUsageTable: dynamodb.TableV2; 21 | constructor(scope: Construct, id: string, props: TenantTokenUsageProps) { 22 | super(scope, id); 23 | 24 | const region = cdk.Stack.of(this).region; 25 | const accountId = cdk.Stack.of(this).account; 26 | 27 | const partitionKey = { 28 | name: "TenantId", 29 | type: dynamodb.AttributeType.STRING, 30 | }; 31 | 32 | // Create the DynamoDB table 33 | const table = new dynamodb.TableV2(this, "TenantTokenUsage", { 34 | tableName: "TenantTokenUsage", 35 | partitionKey: partitionKey, 36 | }); 37 | this.tenantTokenUsageTable = table; 38 | 39 | const tenantTokenUsageCalculatorLambdaExecRole = new iam.Role( 40 | this, 41 | "TenantTokenUsageCalculatorLambdaExecRole", 42 | { 43 | assumedBy: new iam.ServicePrincipal("lambda.amazonaws.com"), 44 | managedPolicies: [ 45 | iam.ManagedPolicy.fromAwsManagedPolicyName( 46 | "CloudWatchLambdaInsightsExecutionRolePolicy" 47 | ), 48 | iam.ManagedPolicy.fromAwsManagedPolicyName( 49 | "service-role/AWSLambdaBasicExecutionRole" 50 | ), 51 | ], 52 | } 53 | ); 54 | 55 | tenantTokenUsageCalculatorLambdaExecRole.addToPolicy( 56 | new iam.PolicyStatement({ 57 | effect: iam.Effect.ALLOW, 58 | actions: ["dynamodb:PutItem"], 59 | resources: [table.tableArn], 60 | }) 61 | ); 62 | 63 | tenantTokenUsageCalculatorLambdaExecRole.addToPolicy( 64 | new iam.PolicyStatement({ 65 | effect: iam.Effect.ALLOW, 66 | actions: [ 67 | "logs:GetQueryResults", 68 | "logs:StartQuery", 69 | "logs:StopQuery", 70 | "logs:FilterLogEvents", 71 | "logs:DescribeLogGroups", 72 | ], 73 | resources: [`arn:aws:logs:${region}:${accountId}:log-group:*`], 74 | }) 75 | ); 76 | 77 | const tenantTokenUsageCalculatorService = new PythonFunction( 78 | this, 79 | "TenantTokenUsageCalculatorService", 80 | { 81 | functionName: "TenantTokenUsageCalculatorService", 82 | entry: path.join(__dirname, "services/tenant-token-usage/"), 83 | runtime: lambda.Runtime.PYTHON_3_12, 84 | index: "tenant_token_usage_calculator.py", 85 | handler: "calculate_daily_tenant_token_usage", 86 | timeout: cdk.Duration.seconds(60), 87 | role: tenantTokenUsageCalculatorLambdaExecRole, 88 | layers: [props.lambdaPowerToolsLayer], 89 | environment: { 90 | TENANT_TOKEN_USAGE_DYNAMODB_TABLE: table.tableName, 91 | }, 92 | } 93 | ); 94 | 95 | const rule = new events.Rule(this, "TenantTokenUsageScheduleRule", { 96 | schedule: events.Schedule.rate(cdk.Duration.minutes(1)), 97 | }); 98 | 99 | // Add the Lambda function as a target of the rule 100 | rule.addTarget( 101 | new targets.LambdaFunction(tenantTokenUsageCalculatorService) 102 | ); 103 | } 104 | } 105 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/usage-plans.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | import { Construct } from "constructs"; 5 | import { 6 | ApiKey, 7 | Period, 8 | type RestApi, 9 | type UsagePlan, 10 | } from "aws-cdk-lib/aws-apigateway"; 11 | 12 | interface UsagePlansProps { 13 | apiGateway: RestApi; 14 | } 15 | 16 | export class UsagePlans extends Construct { 17 | public readonly usagePlanBasicTier: UsagePlan; 18 | constructor(scope: Construct, id: string, props: UsagePlansProps) { 19 | super(scope, id); 20 | 21 | this.usagePlanBasicTier = props.apiGateway.addUsagePlan( 22 | "SaaSGenAIWorkshopUsagePlan", 23 | { 24 | quota: { 25 | limit: 100, 26 | period: Period.DAY, 27 | }, 28 | throttle: { 29 | burstLimit: 30, 30 | rateLimit: 10, 31 | }, 32 | } 33 | ); 34 | 35 | for (const usagePlanTier of [this.usagePlanBasicTier]) { 36 | usagePlanTier.addApiStage({ 37 | api: props.apiGateway, 38 | stage: props.apiGateway.deploymentStage, 39 | }); 40 | } 41 | } 42 | } 43 | -------------------------------------------------------------------------------- /cdk/lib/tenant-template/user-management/user_management_service.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | import boto3 5 | import os 6 | import argparse 7 | 8 | SAAS_APP_USERPOOL_ID=os.environ['SAAS_APP_USERPOOL_ID'] 9 | 10 | cognito = boto3.client('cognito-idp') 11 | 12 | def create_user(tenant_id, email, user_role): 13 | username=email 14 | admin_user_exists = __admin_user_exists(SAAS_APP_USERPOOL_ID,username) 15 | if not admin_user_exists: 16 | temporary_password = "SaaS123!" 17 | response = cognito.admin_create_user( 18 | Username=username, 19 | UserPoolId=SAAS_APP_USERPOOL_ID, 20 | ForceAliasCreation=True, 21 | TemporaryPassword=temporary_password, 22 | MessageAction='SUPPRESS', 23 | UserAttributes=[ 24 | { 25 | 'Name': 'email', 26 | 'Value': email 27 | }, 28 | { 29 | 'Name': 'email_verified', 30 | 'Value': 'true' 31 | }, 32 | { 33 | 'Name': 'custom:tenantId', 34 | 'Value': tenant_id 35 | }, 36 | { 37 | 'Name': 'custom:userRole', 38 | 'Value': user_role 39 | } 40 | ], 41 | ) 42 | 43 | 44 | __set_user_password(SAAS_APP_USERPOOL_ID, username, temporary_password) 45 | else: 46 | response = cognito.admin_update_user_attributes( 47 | UserPoolId=SAAS_APP_USERPOOL_ID, 48 | Username=username, 49 | UserAttributes=[ 50 | { 51 | 'Name': 'email', 52 | 'Value': email 53 | }, 54 | { 55 | 'Name': 'email_verified', 56 | 'Value': 'true' 57 | }, 58 | { 59 | 'Name': 'custom:tenantId', 60 | 'Value': tenant_id 61 | }, 62 | { 63 | 'Name': 'custom:userRole', 64 | 'Value': user_role 65 | } 66 | ] 67 | ) 68 | group_exists = __user_group_exists(SAAS_APP_USERPOOL_ID, tenant_id) 69 | if not group_exists: 70 | __create_user_group(SAAS_APP_USERPOOL_ID, tenant_id) 71 | 72 | __add_user_to_group(SAAS_APP_USERPOOL_ID, username, tenant_id) 73 | return response 74 | 75 | def __admin_user_exists(SAAS_APP_USERPOOL_ID,username): 76 | try: 77 | response=cognito.admin_get_user( 78 | UserPoolId=SAAS_APP_USERPOOL_ID, 79 | Username=username) 80 | return True 81 | except Exception as e: 82 | return False 83 | 84 | def __set_user_password(user_pool_id, username, password): 85 | response = cognito.admin_set_user_password( 86 | UserPoolId=user_pool_id, 87 | Username=username, 88 | Password=password, 89 | Permanent=True 90 | ) 91 | return response 92 | 93 | def __create_user_group(user_pool_id, group_name): 94 | response = cognito.create_group( 95 | GroupName=group_name, 96 | UserPoolId=user_pool_id, 97 | Precedence=0 98 | ) 99 | return response 100 | 101 | 102 | def __add_user_to_group(user_pool_id, user_name, group_name): 103 | response = cognito.admin_add_user_to_group( 104 | UserPoolId=user_pool_id, 105 | Username=user_name, 106 | GroupName=group_name 107 | ) 108 | return response 109 | 110 | def __user_group_exists(user_pool_id, group_name): 111 | try: 112 | response=cognito.get_group( 113 | UserPoolId=user_pool_id, 114 | GroupName=group_name) 115 | return True 116 | except Exception as e: 117 | return False 118 | 119 | if __name__ == '__main__': 120 | parser = argparse.ArgumentParser( 121 | description='Managing tenant users') 122 | parser.add_argument('--tenant-id', type=str, help='tenant-id', required=True) 123 | parser.add_argument('--email', type=str, help='email', required=True) 124 | parser.add_argument('--user-role', type=str, help='user role', required=True) 125 | 126 | args = parser.parse_args() 127 | create_user(**vars(args)) 128 | -------------------------------------------------------------------------------- /cdk/package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "saas-genai-workshop", 3 | "version": "0.1.0", 4 | "bin": { 5 | "saas-genai-workshop": "bin/saas-genai-workshop.js" 6 | }, 7 | "scripts": { 8 | "build": "tsc", 9 | "watch": "tsc -w", 10 | "test": "jest", 11 | "cdk": "cdk" 12 | }, 13 | "devDependencies": { 14 | "@types/jest": "^29.5.11", 15 | "@types/node": "20.10.4", 16 | "aws-cdk": "2.177.0", 17 | "jest": "^29.7.0", 18 | "ts-jest": "^29.1.1", 19 | "ts-node": "^10.9.2", 20 | "typescript": "~5.3.3" 21 | }, 22 | "dependencies": { 23 | "@cdklabs/generative-ai-cdk-constructs": "0.1.106", 24 | "@cdklabs/sbt-aws": "^0.3.6", 25 | "aws-cdk-lib": "^2.177.0", 26 | "constructs": "^10.0.0", 27 | "source-map-support": "^0.5.21" 28 | } 29 | } 30 | -------------------------------------------------------------------------------- /cdk/test/cdk.test.ts: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: MIT-0 3 | 4 | // import * as cdk from 'aws-cdk-lib'; 5 | // import { Template } from 'aws-cdk-lib/assertions'; 6 | // import * as Cdk from '../lib/cdk-stack'; 7 | 8 | // example test. To run these tests, uncomment this file along with the 9 | // example resource in lib/cdk-stack.ts 10 | test("SQS Queue Created", () => { 11 | // const app = new cdk.App(); 12 | // // WHEN 13 | // const stack = new Cdk.CdkStack(app, 'MyTestStack'); 14 | // // THEN 15 | // const template = Template.fromStack(stack); 16 | // template.hasResourceProperties('AWS::SQS::Queue', { 17 | // VisibilityTimeout: 300 18 | // }); 19 | }); 20 | -------------------------------------------------------------------------------- /cdk/tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "target": "ES2020", 4 | "module": "commonjs", 5 | "lib": [ 6 | "es2020", 7 | "dom" 8 | ], 9 | "declaration": true, 10 | "strict": true, 11 | "noImplicitAny": true, 12 | "strictNullChecks": true, 13 | "noImplicitThis": true, 14 | "alwaysStrict": true, 15 | "noUnusedLocals": false, 16 | "noUnusedParameters": false, 17 | "noImplicitReturns": true, 18 | "noFallthroughCasesInSwitch": false, 19 | "inlineSourceMap": true, 20 | "inlineSources": true, 21 | "experimentalDecorators": true, 22 | "strictPropertyInitialization": false, 23 | "typeRoots": [ 24 | "./node_modules/@types" 25 | ] 26 | }, 27 | "exclude": [ 28 | "node_modules", 29 | "cdk.out" 30 | ] 31 | } 32 | -------------------------------------------------------------------------------- /data/cur_report.zip: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/aws-saas-genai-rag-workshop/e81b3a60cdb6adaf4d1f9c1fb11e279b279d473d/data/cur_report.zip -------------------------------------------------------------------------------- /data/tenant1-meeting-notes.txt: -------------------------------------------------------------------------------- 1 | 2 | Meeting Notes for Mining Companies - North America 3 | 4 | 1. Meeting Date: 2024-09-01 5 | Sites Discussed: Yukon Gold Site, Alaska 6 | Activities: Initial drilling for gold reserves commenced. Preliminary tests show promising yields. Environmental impact assessments completed and sent to regulatory bodies. 7 | 8 | 2. Meeting Date: 2024-09-03 9 | Sites Discussed: Copper Ridge Site, Arizona 10 | Activities: Surface extraction is in full swing. Copper purity levels are being tested, with initial results showing higher than expected concentrations. Equipment maintenance is scheduled for next quarter. 11 | 12 | 3. Meeting Date: 2024-09-05 13 | Sites Discussed: Ironclad Site, Minnesota 14 | Activities: Workforce expansion plans discussed due to increased demand for iron ore. Safety protocols for deep mining were reviewed following a minor accident last month. 15 | 16 | 4. Meeting Date: 2024-09-07 17 | Sites Discussed: Thunder Bay Lithium Site, Ontario 18 | Activities: Phase 2 of the extraction plan initiated. The geologists are mapping new veins discovered last week. Water usage strategies to be revised due to concerns from local communities. 19 | 20 | 5. Meeting Date: 2024-09-09 21 | Sites Discussed: Diamond Valley Site, Northwest Territories 22 | Activities: Security upgrades for diamond shipments were approved. Extraction levels are steady, and new diamond-cutting equipment will be delivered next month. 23 | 24 | 6. Meeting Date: 2024-09-11 25 | Sites Discussed: Coal Mountain Site, Wyoming 26 | Activities: Reclamation plans discussed as the site approaches end-of-life. Discussions about transitioning some workforce to nearby new exploration projects. Environmental rehabilitation measures will begin in 2025. 27 | 28 | 7. Meeting Date: 2024-09-13 29 | Sites Discussed: Silver Creek Site, Nevada 30 | Activities: Drilling continues, with silver output slightly below projections. Additional geological surveys were approved to investigate deeper veins. Team noted delays in machinery due to supply chain issues. 31 | 32 | 8. Meeting Date: 2024-09-15 33 | Sites Discussed: Zinc Lake Site, Manitoba 34 | Activities: Safety audit results were shared with the team, highlighting areas for improvement. New ventilation systems are to be installed by year-end. Production rates remain stable. 35 | 36 | 9. Meeting Date: 2024-09-17 37 | Sites Discussed: Uranium Peak Site, New Mexico 38 | Activities: New radiation monitoring protocols implemented. Export permits for uranium are being fast-tracked with government negotiations ongoing. Expansion into neighboring areas is under consideration. 39 | 40 | 10. Meeting Date: 2024-09-19 41 | Sites Discussed: Cobalt Hills Site, Idaho 42 | Activities: Cobalt extraction operations increased to meet rising demand for battery materials. A new partnership with a local transportation firm was secured to ensure timely deliveries to refining facilities. 43 | 44 | End of Meeting Notes. 45 | -------------------------------------------------------------------------------- /data/tenant2-meeting-notes.txt: -------------------------------------------------------------------------------- 1 | 2 | Meeting Notes for Financial Bank Companies - Next Year Plan 3 | 4 | 1. Meeting Date: 2024-09-01 5 | Products Discussed: Personal Loans, Mortgages 6 | Plan: Introduce new flexible payment options for personal loans. Mortgage interest rates to be reviewed for possible reductions in Q1 to attract more homebuyers. 7 | 8 | 2. Meeting Date: 2024-09-03 9 | Products Discussed: Credit Cards, Savings Accounts 10 | Plan: Expand rewards program for credit card users, focusing on travel and dining benefits. New high-yield savings account with competitive rates will be launched by Q2. 11 | 12 | 3. Meeting Date: 2024-09-05 13 | Products Discussed: Investment Portfolios, Retirement Plans 14 | Plan: Revamp investment portfolio offerings with a focus on sustainable and green investments. A new marketing campaign for retirement plans targeting millennials will roll out in Q3. 15 | 16 | 4. Meeting Date: 2024-09-07 17 | Products Discussed: Business Loans, SME Financing 18 | Plan: Increase lending capacity for small and medium enterprises. Develop specialized loan packages for startups, with reduced interest rates and extended repayment terms. 19 | 20 | 5. Meeting Date: 2024-09-09 21 | Products Discussed: Mobile Banking, Online Payment Services 22 | Plan: Upgrade mobile banking app with more user-friendly features and enhanced security. Expand partnerships with online payment platforms to improve transaction speed and reduce fees. 23 | 24 | 6. Meeting Date: 2024-09-11 25 | Products Discussed: Wealth Management, Private Banking 26 | Plan: Introduce personalized wealth management services using AI-driven insights. Expand private banking services to include exclusive investment opportunities for high-net-worth clients. 27 | 28 | 7. Meeting Date: 2024-09-13 29 | Products Discussed: Insurance Products, Health Coverage 30 | Plan: Roll out new health insurance packages with broader coverage. Collaborate with healthcare providers to offer bundled services and discounts for insured customers. 31 | 32 | 8. Meeting Date: 2024-09-15 33 | Products Discussed: Student Loans, Educational Savings Plans 34 | Plan: Adjust student loan interest rates in line with government policies. Launch a new educational savings plan with tax incentives for parents and guardians. 35 | 36 | 9. Meeting Date: 2024-09-17 37 | Products Discussed: Corporate Bonds, Treasury Services 38 | Plan: Expand corporate bond offerings, targeting high-growth industries. Upgrade treasury management services to include more automation and risk management tools. 39 | 40 | 10. Meeting Date: 2024-09-19 41 | Products Discussed: Foreign Exchange Services, International Transfers 42 | Plan: Launch a new foreign exchange platform with lower transaction fees for cross-border payments. Increase market presence in Asia and Latin America through strategic partnerships. 43 | 44 | End of Meeting Notes. 45 | -------------------------------------------------------------------------------- /scripts/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/aws-saas-genai-rag-workshop/e81b3a60cdb6adaf4d1f9c1fb11e279b279d473d/scripts/.DS_Store -------------------------------------------------------------------------------- /scripts/cleanup.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash -e 2 | 3 | export CDK_PARAM_SYSTEM_ADMIN_EMAIL="$1" 4 | 5 | if [[ -z "$CDK_PARAM_SYSTEM_ADMIN_EMAIL" ]]; then 6 | echo "Please provide system admin email" 7 | exit 1 8 | fi 9 | 10 | echo "$(date) emptying out buckets..." 11 | for i in $(aws s3 ls | awk '{print $3}' | grep -E "^saas-genai-workshop-*"); do 12 | echo "$(date) emptying out s3 bucket with name s3://${i}..." 13 | aws s3 rm --recursive "s3://${i}" 14 | done 15 | 16 | # Function to check if a command exists 17 | command_exists() { 18 | command -v "$1" >/dev/null 2>&1 19 | } 20 | 21 | # Check if required tools are installed 22 | if ! command_exists aws; then 23 | echo "Error: AWS CLI is not installed. Please install it and configure your credentials." 24 | exit 1 25 | fi 26 | 27 | if ! command_exists cdk; then 28 | echo "Error: AWS CDK is not installed. Please install it using 'npm install -g aws-cdk'." 29 | exit 1 30 | fi 31 | 32 | # Navigate to the directory containing the CDK app 33 | # Replace 'path/to/cdk/app' with the actual path to your CDK application 34 | cd ../cdk 35 | 36 | echo "Starting cleanup process..." 37 | 38 | # Destroy the CDK stacks 39 | echo "Destroying CDK stacks..." 40 | cdk destroy --all --force 41 | 42 | # Remove the CDK context file 43 | echo "Removing CDK context file..." 44 | rm -f cdk.context.json 45 | 46 | # Remove the CDK output directory 47 | echo "Removing CDK output directory..." 48 | rm -rf cdk.out 49 | 50 | # Optionally, remove node_modules if you want to clean up dependencies 51 | # echo "Removing node_modules..." 52 | # rm -rf node_modules 53 | 54 | # Clean up any DynamoDB tables 55 | echo "Cleaning up DynamoDB tables..." 56 | tables=$(aws dynamodb list-tables --query 'TableNames[?contains(@, `TenantCostAndUsageAttribution`)]' --output text) 57 | 58 | for table in $tables; do 59 | echo "Deleting DynamoDB table: $table" 60 | aws dynamodb delete-table --table-name $table --no-cli-pager 61 | done 62 | 63 | # Clean up any Lambda functions 64 | # echo "Cleaning up Lambda functions..." 65 | # functions=$(aws lambda list-functions --query 'Functions[?contains(FunctionName, `YourProjectPrefix`)].FunctionName' --output text) 66 | 67 | # for func in $functions; do 68 | # echo "Deleting Lambda function: $func" 69 | # aws lambda delete-function --function-name $func 70 | # done 71 | 72 | # Clean up CodeCommit repository 73 | echo "Cleaning up CodeCommit repository..." 74 | export CDK_PARAM_CODE_COMMIT_REPOSITORY_NAME="saas-genai-workshop" 75 | 76 | if aws codecommit get-repository --repository-name $CDK_PARAM_CODE_COMMIT_REPOSITORY_NAME >/dev/null 2>&1; then 77 | echo "Deleting CodeCommit repository: $CDK_PARAM_CODE_COMMIT_REPOSITORY_NAME" 78 | aws codecommit delete-repository --repository-name $CDK_PARAM_CODE_COMMIT_REPOSITORY_NAME --no-cli-pager 79 | echo "CodeCommit repository deleted successfully." 80 | else 81 | echo "CodeCommit repository $CDK_PARAM_CODE_COMMIT_REPOSITORY_NAME does not exist. Skipping deletion." 82 | fi 83 | 84 | echo "$(date) cleaning up log groups..." 85 | next_token="" 86 | while true; do 87 | if [[ "${next_token}" == "" ]]; then 88 | response=$(aws logs describe-log-groups) 89 | else 90 | response=$(aws logs describe-log-groups --starting-token "$next_token") 91 | fi 92 | 93 | log_groups=$(echo "$response" | jq -r '.logGroups[].logGroupName | select(. | test("^/aws/lambda/ControlPlaneStack-*|saas-genai-workshop-*|provisioningJobRunner*|^/aws/lambda/TenantCostCalculatorService|^/aws/bedrock/SaaSGenAIWorkshopBedrockLogGroup"))') 94 | for i in $log_groups; do 95 | if [[ -z "${skip_flag}" ]]; then 96 | read -p "Delete log group with name $i [Y/n] " -n 1 -r 97 | fi 98 | 99 | if [[ $REPLY =~ ^[n]$ ]]; then 100 | echo "$(date) NOT deleting log group $i." 101 | else 102 | echo "$(date) deleting log group with name $i..." 103 | aws logs delete-log-group --log-group-name "$i" 104 | fi 105 | done 106 | 107 | next_token=$(echo "$response" | jq '.NextToken') 108 | if [[ "${next_token}" == "null" ]]; then 109 | # no more results left. Exit loop... 110 | break 111 | fi 112 | done 113 | 114 | echo "$(date) cleaning up user pools..." 115 | next_token="" 116 | while true; do 117 | if [[ "${next_token}" == "" ]]; then 118 | response=$( aws cognito-idp list-user-pools --max-results 10) 119 | else 120 | # using next-token instead of starting-token. See: https://github.com/aws/aws-cli/issues/7661 121 | response=$( aws cognito-idp list-user-pools --max-results 10 --next-token "$next_token") 122 | fi 123 | 124 | pool_ids=$(echo "$response" | jq -r '.UserPools[] | select(.Name | test("^IdentityProvidertenantUserPool|^CognitoAuthUserPool")) |.Id') 125 | for i in $pool_ids; do 126 | echo "$(date) deleting user pool with name $i..." 127 | echo "getting pool domain..." 128 | pool_domain=$(aws cognito-idp describe-user-pool --user-pool-id "$i" | jq -r '.UserPool.Domain') 129 | 130 | # Delete the pool domain if it exists 131 | if [[ "$pool_domain" != "null" && -n "$pool_domain" ]]; then 132 | echo "deleting pool domain $pool_domain..." 133 | aws cognito-idp delete-user-pool-domain \ 134 | --user-pool-id "$i" \ 135 | --domain "$pool_domain" 136 | else 137 | echo "No domain associated with this user pool or unable to retrieve domain." 138 | fi 139 | 140 | echo "deleting pool $i..." 141 | aws cognito-idp delete-user-pool --user-pool-id "$i" --no-cli-pager 142 | done 143 | 144 | next_token=$(echo "$response" | jq -r '.NextToken') 145 | if [[ "${next_token}" == "null" ]]; then 146 | # no more results left. Exit loop... 147 | break 148 | fi 149 | done 150 | 151 | echo "Cleanup process completed." 152 | -------------------------------------------------------------------------------- /scripts/deploy.sh: -------------------------------------------------------------------------------- 1 | # Check if stack name parameter is provided 2 | if [ $# -ne 1 ]; then 3 | echo "Usage: $0 " 4 | echo "Example: $0 my-stack-name" 5 | exit 1 6 | fi 7 | 8 | # Get the stack name from command line argument 9 | STACK_NAME="$1" 10 | 11 | cd ../cdk 12 | npx cdk deploy $STACK_NAME --require-approval never --concurrency 10 --asset-parallelism true --exclusively 13 | 14 | # Upload the updated code to the S3 bucket 15 | S3_TENANT_SOURCECODE_BUCKET_URL=$(aws cloudformation describe-stacks --stack-name saas-genai-workshop-bootstrap-template --query "Stacks[0].Outputs[?OutputKey=='TenantSourceCodeS3Bucket'].OutputValue" --output text) 16 | echo "S3 bucket url: $S3_TENANT_SOURCECODE_BUCKET_URL" 17 | 18 | cd .. 19 | echo "Uploading updated code...." 20 | aws s3 sync "." "s3://$S3_TENANT_SOURCECODE_BUCKET_URL" --exclude "cdk/cdk.out/*" --exclude "cdk/node_modules/*" --exclude ".git/*" 21 | echo "Completed uploading updated code." 22 | 23 | 24 | 25 | -------------------------------------------------------------------------------- /scripts/get_opensearch_indices.py: -------------------------------------------------------------------------------- 1 | import boto3 2 | import json 3 | import time 4 | import os 5 | from opensearchpy import OpenSearch, RequestsHttpConnection 6 | from requests_aws4auth import AWS4Auth 7 | 8 | # Constants 9 | PORT = 443 10 | STACK_NAME = "saas-genai-workshop-bootstrap-template" 11 | OSSC_ARN_PARAM_NAME = "SaaSGenAIWorkshopOSSCollectionArn" 12 | 13 | # AWS configuration 14 | region = os.environ.get('AWS_REGION', 'us-west-2') 15 | service = 'aoss' 16 | credentials = boto3.Session().get_credentials() 17 | awsauth = AWS4Auth(credentials.access_key, credentials.secret_key, region, service, session_token=credentials.token) 18 | 19 | 20 | def get_collection_arn(): 21 | # OpenSearch Serverless configuration 22 | cf_client = boto3.Session(region_name=region).client('cloudformation') 23 | # Get the ARN using list comprehension 24 | oass_collection_arn = next( 25 | (output['OutputValue'] 26 | for output in cf_client.describe_stacks(StackName=STACK_NAME)['Stacks'][0]['Outputs'] 27 | if output['OutputKey'] == OSSC_ARN_PARAM_NAME), 28 | None 29 | ) 30 | return (oass_collection_arn) 31 | 32 | aoss_client = boto3.Session(region_name=region).client('opensearchserverless') 33 | 34 | def __get_current_role(): 35 | # Get the full STS ARN 36 | sts_arn = boto3.client('sts').get_caller_identity()['Arn'] 37 | 38 | # Extract the role name using string operations 39 | role_name = sts_arn.split('/')[1] # Get the middle part between first and last '/' 40 | 41 | # Format it as IAM role ARN 42 | account_id = sts_arn.split(':')[4] # Get the AWS account ID 43 | iam_role_arn = f"arn:aws:iam::{account_id}:role/{role_name}" 44 | 45 | return(iam_role_arn) 46 | 47 | # Function creating a Data Access Policy per tenant 48 | def __add_data_access_policy(kb_collection): 49 | kb_collection_name=kb_collection['name'] 50 | iam_role_arn = __get_current_role() 51 | try: 52 | response = aoss_client.create_access_policy( 53 | name=f'participantpolicy', 54 | description=f'Data access policy for participant for collection: {kb_collection_name}', 55 | policy=json.dumps(__generate_data_access_policy(iam_role_arn, kb_collection_name)), 56 | type='data') 57 | 58 | print(f'Participant data access policy created') 59 | except boto3.exceptions.botocore.exceptions.ClientError as e: 60 | error_code = e.response['Error']['Code'] 61 | if error_code == 'ConflictException': 62 | print(f'Policy with name "participantpolicy" and type "data" already exists. Skipping creation.') 63 | else: 64 | print('Error occurred while adding data access policy', e) 65 | raise Exception('Error occurred while adding data access policy') from e 66 | 67 | # Generating a Data Access Policy per tenant 68 | def __generate_data_access_policy(iam_role_arn, kb_collection_name): 69 | return [ 70 | { 71 | "Rules": [ 72 | { 73 | "Resource": [ 74 | f"index/{kb_collection_name}/*" 75 | ], 76 | "Permission": [ 77 | "aoss:DescribeIndex" 78 | ], 79 | "ResourceType": "index" 80 | } 81 | ], 82 | "Principal": [iam_role_arn], 83 | } 84 | ] 85 | 86 | def __get_opensearch_serverless_collection_details(collection_arn): 87 | try: 88 | kb_collection_id = collection_arn.split('/')[-1] 89 | kb_collection = aoss_client.batch_get_collection( 90 | ids=[kb_collection_id] 91 | ) 92 | 93 | kb_collection_endpoint = kb_collection['collectionDetails'][0]['collectionEndpoint'] 94 | return kb_collection['collectionDetails'][0] 95 | except Exception as e: 96 | print('Error occured while getting OpenSearch serverless collection details', e) 97 | raise Exception('Error occured while getting OpenSearch serverless collection details') from e 98 | 99 | def get_index_sizes(kb_collection): 100 | kb_collection_endpoint=kb_collection['collectionEndpoint'] 101 | kb_collection_endpoint_domain= kb_collection_endpoint.split("//")[-1] 102 | # Create the OpenSearch client 103 | client = OpenSearch( 104 | hosts=[{'host': kb_collection_endpoint_domain, 'port': PORT}], 105 | http_auth=awsauth, 106 | use_ssl=True, 107 | verify_certs=True, 108 | connection_class=RequestsHttpConnection 109 | ) 110 | 111 | # Get all indexes in the collection 112 | indices = client.cat.indices(format='json') 113 | 114 | ## Print index name and data size 115 | for index in indices: 116 | index_name = index['index'] 117 | store_size = index['store.size'] 118 | print(f"Index: {index_name}, Data Size: {store_size}") 119 | 120 | if __name__ == "__main__": 121 | collection_arn=get_collection_arn() 122 | kb_collection=__get_opensearch_serverless_collection_details(collection_arn) 123 | __add_data_access_policy(kb_collection) 124 | time.sleep(10) 125 | get_index_sizes(kb_collection) -------------------------------------------------------------------------------- /scripts/install.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash -e 2 | 3 | export CDK_PARAM_SYSTEM_ADMIN_EMAIL="$1" 4 | 5 | if [[ -z "$CDK_PARAM_SYSTEM_ADMIN_EMAIL" ]]; then 6 | echo "Please provide system admin email" 7 | exit 1 8 | fi 9 | 10 | # Check if running on EC2 by looking for the AWS_REGION environment variable 11 | if [[ -n "$AWS_REGION" ]]; then 12 | REGION="$AWS_REGION" 13 | else 14 | # If not on EC2, try to get the region from aws configure 15 | REGION=$(aws configure get region) 16 | if [[ -z "$REGION" ]]; then 17 | echo "Unable to determine AWS region. Please set AWS_REGION environment variable or configure AWS CLI." 18 | exit 1 19 | fi 20 | fi 21 | 22 | # Preprovision base infrastructure 23 | cd ../cdk 24 | npm install 25 | 26 | npx cdk bootstrap 27 | npx cdk deploy --all --require-approval never --concurrency 10 --asset-parallelism true 28 | 29 | CP_API_GATEWAY_URL=$(aws cloudformation describe-stacks --stack-name ControlPlaneStack --query "Stacks[0].Outputs[?OutputKey=='controlPlaneAPIEndpoint'].OutputValue" --output text) 30 | 31 | echo "Control plane api gateway url: $CP_API_GATEWAY_URL" 32 | 33 | 34 | S3_TENANT_SOURCECODE_BUCKET_URL=$(aws cloudformation describe-stacks --stack-name saas-genai-workshop-bootstrap-template --query "Stacks[0].Outputs[?OutputKey=='TenantSourceCodeS3Bucket'].OutputValue" --output text) 35 | echo "S3 bucket url: $S3_TENANT_SOURCECODE_BUCKET_URL" 36 | 37 | # Step 3: Define folder to upload and target S3 bucket 38 | SCRIPT_DIR="$(dirname "$(realpath "$0")")" # Get the directory of the install.sh script 39 | FOLDER_PATH="$(dirname "$SCRIPT_DIR")" # Get the parent folder of the script 40 | 41 | # Step 4: Upload the folder to the S3 bucket 42 | echo "Uploading folder $FOLDER_PATH to S3 $S3_TENANT_SOURCECODE_BUCKET_URL" 43 | aws s3 cp "$FOLDER_PATH" "s3://$S3_TENANT_SOURCECODE_BUCKET_URL" --recursive --exclude "cdk/cdk.out/*" --exclude "cdk/node_modules/*" --exclude ".git/*" --quiet 44 | 45 | echo "Installation completed successfully" -------------------------------------------------------------------------------- /scripts/labs/lab1_add_tenant_provision_service.py: -------------------------------------------------------------------------------- 1 | from replace_code import replace_code_in_file 2 | 3 | original_str = """# TODO: Lab1 - Add tenant provisioning service""" 4 | 5 | update_str = """# TODO: Lab1 - Add tenant provisioning service 6 | tenant_provision_output=$(python3 lib/tenant-template/tenant-provisioning/tenant_provisioning_service.py --tenantid "$CDK_PARAM_TENANT_ID" 2>&1 > /dev/null && exit_code=$?) || exit_code=$? 7 | check_error "$provision_name" $exit_code "$tenant_provision_output" 8 | """ 9 | 10 | replace_code_in_file("../provision-tenant.sh", original_str, update_str) 11 | -------------------------------------------------------------------------------- /scripts/labs/lab1_provision_tenant_resources.py: -------------------------------------------------------------------------------- 1 | from replace_code import replace_code_in_file 2 | 3 | original_str = """# TODO: Lab1 - Add provision tenant resources""" 4 | 5 | update_str = """# TODO: Lab1 - Add provision tenant resources 6 | __create_opensearch_serverless_tenant_index(tenant_id, kb_collection_endpoint_domain) 7 | __create_s3_tenant_prefix(tenant_id, rule_name) 8 | __create_tenant_knowledge_base(tenant_id, kb_collection_name, rule_name) 9 | """ 10 | 11 | replace_code_in_file("../../cdk/lib/tenant-template/tenant-provisioning/tenant_provisioning_service.py", original_str, update_str) 12 | -------------------------------------------------------------------------------- /scripts/labs/lab2_add_principal_tag.py: -------------------------------------------------------------------------------- 1 | from replace_code import replace_code_in_file 2 | 3 | original_str = """// TODO: Lab2 - Add principalTag in ABAC policy 4 | resourceName: "*",""" 5 | 6 | update_str = """// TODO: Lab2 - Add principalTag in ABAC policy 7 | resourceName: "${aws:PrincipalTag/KnowledgeBaseId}",""" 8 | 9 | replace_code_in_file("../../cdk/lib/tenant-template/services.ts", original_str, update_str) 10 | 11 | 12 | -------------------------------------------------------------------------------- /scripts/labs/lab2_hardcode_knowledgebase_id.py: -------------------------------------------------------------------------------- 1 | from replace_code import replace_code_in_file 2 | 3 | def replace_kb_id(kb_id): 4 | 5 | original_str = """# TODO: Lab2 - uncomment below and hardcode an knowledge base id 6 | # knowledge_base_id = "" 7 | # logger.info(f"hard coded knowledge base id: {knowledge_base_id}")""" 8 | 9 | update_str = f"""# TODO: Lab2 - uncomment below and hardcode an knowledge base id 10 | knowledge_base_id = "{kb_id}" 11 | logger.info(f"hard coded knowledge base id: {{knowledge_base_id}}") 12 | """ 13 | 14 | replace_code_in_file("../../cdk/lib/tenant-template/services/ragService/rag_service.py", original_str, update_str) 15 | 16 | 17 | if __name__ == "__main__": 18 | kb_id = input("Enter the knowledge base id: ") 19 | replace_kb_id(kb_id) 20 | -------------------------------------------------------------------------------- /scripts/labs/lab2_remove_hardcode_knowledgebase_id.py: -------------------------------------------------------------------------------- 1 | from replace_code import replace_code_in_file_regex 2 | 3 | old_pattern = r"""# TODO: Lab2 - uncomment below and hardcode an knowledge base id 4 | \s*knowledge_base_id = ".*?" 5 | \s*logger\.info\(f"hard coded knowledge base id: \{knowledge_base_id\}"\)""" 6 | 7 | update_str = """# TODO: Lab2 - uncomment below and hardcode an knowledge base id 8 | # knowledge_base_id = "" 9 | # logger.info(f"hard coded knowledge base id: {knowledge_base_id}") 10 | """ 11 | 12 | replace_code_in_file_regex("../../cdk/lib/tenant-template/services/ragService/rag_service.py", old_pattern, update_str) 13 | 14 | 15 | -------------------------------------------------------------------------------- /scripts/labs/lab3_enable_tenant_token_usage.py: -------------------------------------------------------------------------------- 1 | from replace_code import replace_code_in_file 2 | 3 | original_str = """# TODO: Lab3 - Enable tenant token usage""" 4 | 5 | update_str = """# TODO: Lab3 - Enable tenant token usage 6 | if ('/invoke' in method_arn and __is_tenant_token_limit_exceeded(tenant_id, input_tokens, output_tokens)) : 7 | return authorizer_layer.create_auth_denied_policy(method_arn) 8 | """ 9 | 10 | replace_code_in_file("../../cdk/lib/tenant-template/services/authorizerService/tenant_authorizer.py", original_str, update_str) 11 | -------------------------------------------------------------------------------- /scripts/labs/lab4_calculate_input_outputs_tokens_attribution.py: -------------------------------------------------------------------------------- 1 | from replace_code import replace_code_in_file 2 | 3 | original_input_output_tokens_str = """# TODO: Lab4 - Add Amazon CloudWatch logs insights queries for converse input output tokens 4 | converse_input_output_tokens_query = "" 5 | """ 6 | 7 | update_input_output_tokens_str = r"""# TODO: Lab4 - Add Amazon CloudWatch logs insights queries for converse input output tokens 8 | converse_input_output_tokens_query = "filter @message like /ModelInvocationInputTokens|ModelInvocationOutputTokens/ \ 9 | | fields tenant_id as TenantId, ModelInvocationInputTokens.0 as ModelInvocationInputTokens, ModelInvocationOutputTokens.0 as ModelInvocationOutputTokens \ 10 | | stats sum(ModelInvocationInputTokens) as TotalInputTokens, sum(ModelInvocationOutputTokens) as TotalOutputTokens by TenantId, dateceil(@timestamp, 1d) as timestamp" 11 | """ 12 | 13 | replace_code_in_file("../../cdk/lib/tenant-template/services/aggregate-metrics/invoke_model_tenant_cost.py", original_input_output_tokens_str, update_input_output_tokens_str) 14 | 15 | original_total_input_output_tokens_str = """# TODO: Lab4 - Add Amazon CloudWatch logs insights queries to get total converse input output tokens 16 | total_converse_input_output_tokens_query = "" 17 | """ 18 | 19 | update_total_input_output_tokens_str = r"""# TODO: Lab4 - Add Amazon CloudWatch logs insights queries to get total converse input output tokens 20 | total_converse_input_output_tokens_query = "filter @message like /ModelInvocationInputTokens|ModelInvocationOutputTokens/ \ 21 | | fields ModelInvocationInputTokens.0 as ModelInvocationInputTokens, ModelInvocationOutputTokens.0 as ModelInvocationOutputTokens \ 22 | | stats sum(ModelInvocationInputTokens) as TotalInputTokens, sum(ModelInvocationOutputTokens) as TotalOutputTokens by dateceil(@timestamp, 1d) as timestamp" 23 | """ 24 | 25 | replace_code_in_file("../../cdk/lib/tenant-template/services/aggregate-metrics/invoke_model_tenant_cost.py", original_total_input_output_tokens_str, update_total_input_output_tokens_str) 26 | 27 | 28 | original_input_output_tokens_attribution_str = """# TODO: Lab4 - Calculate the percentage of tenant attribution for converse input and output tokens 29 | tenant_attribution_input_tokens_percentage = 0 30 | tenant_attribution_output_tokens_percentage = 0 31 | """ 32 | 33 | update_input_output_tokens_attribution_str = """# TODO: Lab4 - Calculate the percentage of tenant attribution for converse input and output tokens 34 | tenant_attribution_input_tokens_percentage = tenant_input_tokens/total_input_tokens 35 | tenant_attribution_output_tokens_percentage = tenant_output_tokens/total_input_tokens 36 | """ 37 | 38 | replace_code_in_file("../../cdk/lib/tenant-template/services/aggregate-metrics/invoke_model_tenant_cost.py", original_input_output_tokens_attribution_str, update_input_output_tokens_attribution_str) 39 | 40 | -------------------------------------------------------------------------------- /scripts/labs/lab4_calculate_kb_input_tokens_attribution.py: -------------------------------------------------------------------------------- 1 | from replace_code import replace_code_in_file 2 | 3 | original_kb_tokens_str = """#TODO: Lab4 - Add Amazon CloudWatch logs insights queries to get knowledge base input tokens 4 | knowledgebase_input_tokens_query = "" 5 | """ 6 | 7 | update_kb_tokens_str = r"""#TODO: Lab4 - Add Amazon CloudWatch logs insights queries to get knowledge base input tokens 8 | knowledgebase_input_tokens_query = "fields @timestamp, identity.arn, input.inputTokenCount \ 9 | | filter modelId like /amazon.titan-embed-text-v1/ and operation = 'InvokeModel' \ 10 | | parse identity.arn '/bedrock-kb-role-*/' as tenantId \ 11 | | filter ispresent(tenantId) \ 12 | | stats sum(input.inputTokenCount) as TotalInputTokens by tenantId, dateceil(@timestamp, 1d) as timestamp \ 13 | | sort totalInputTokenCount desc" 14 | """ 15 | 16 | replace_code_in_file("../../cdk/lib/tenant-template/services/aggregate-metrics/invoke_model_tenant_cost.py", original_kb_tokens_str, update_kb_tokens_str) 17 | 18 | original_total_kb_tokens_str = """# TODO: Lab4 - Add Amazon CloudWatch logs insights queries to get total knowledge base input tokens 19 | total_knowledgebase_input_tokens_query = "" 20 | """ 21 | 22 | update_total_kb_tokens_str = r"""# TODO: Lab4 - Add Amazon CloudWatch logs insights queries to get total knowledge base input tokens 23 | total_knowledgebase_input_tokens_query = "fields @timestamp, identity.arn, input.inputTokenCount \ 24 | | filter modelId like /amazon.titan-embed-text-v1/ and operation = 'InvokeModel' \ 25 | | parse identity.arn '/bedrock-kb-role-*/' as tenantId \ 26 | | filter ispresent(tenantId) \ 27 | | stats sum(input.inputTokenCount) as TotalInputTokens, dateceil(@timestamp, 1d) as timestamp" 28 | """ 29 | 30 | replace_code_in_file("../../cdk/lib/tenant-template/services/aggregate-metrics/invoke_model_tenant_cost.py", original_total_kb_tokens_str, update_total_kb_tokens_str) 31 | 32 | 33 | original_kb_tokens_attribution_str = """# TODO: Lab4 - Calculate the percentage of tenant attribution for knowledge base input tokens 34 | tenant_kb_input_tokens_attribution_percentage = 0 35 | """ 36 | 37 | update_kb_tokens_attribution_str = """# TODO: Lab4 - Calculate the percentage of tenant attribution for knowledge base input tokens 38 | tenant_kb_input_tokens_attribution_percentage = input_tokens/total_knowledgebase_input_tokens 39 | """ 40 | 41 | replace_code_in_file("../../cdk/lib/tenant-template/services/aggregate-metrics/invoke_model_tenant_cost.py", original_kb_tokens_attribution_str, update_kb_tokens_attribution_str) 42 | 43 | -------------------------------------------------------------------------------- /scripts/labs/lab4_calculate_tenant_cost_kb.py: -------------------------------------------------------------------------------- 1 | from replace_code import replace_code_in_file 2 | 3 | original_tenant_cost_kb_str = """# TODO: Lab4 - Calculate tenant cost for ingesting & retrieving tenant data to/from Amazon Bedrock Knowledge Base 4 | tenant_kb_input_tokens_cost = 0 5 | """ 6 | 7 | update_tenant_cost_kb_str = """# TODO: Lab4 - Calculate tenant cost for ingesting & retrieving tenant data to/from Amazon Bedrock Knowledge Base 8 | tenant_kb_input_tokens_cost = self.__get_tenant_cost(EMBEDDING_TITAN_INPUT_TOKENS_LABEL, total_service_cost_dict, tenant_attribution_percentage_json) 9 | """ 10 | 11 | replace_code_in_file("../../cdk/lib/tenant-template/services/aggregate-metrics/invoke_model_tenant_cost.py", original_tenant_cost_kb_str, update_tenant_cost_kb_str) 12 | 13 | -------------------------------------------------------------------------------- /scripts/labs/lab4_calculate_tenant_cost_service.py: -------------------------------------------------------------------------------- 1 | from replace_code import replace_code_in_file 2 | 3 | original_tenant_cost_service_str = """# TODO: Lab4 - Calculate tenant cost for generating final tenant specific response 4 | tenant_input_tokens_cost = 0 5 | tenant_output_tokens_cost = 0 6 | """ 7 | 8 | update_tenant_cost_service_str = """# TODO: Lab4 - Calculate tenant cost for generating final tenant specific response 9 | tenant_input_tokens_cost = self.__get_tenant_cost(TEXTLITE_INPUT_TOKENS_LABEL, total_service_cost_dict, tenant_attribution_percentage_json) 10 | tenant_output_tokens_cost = self.__get_tenant_cost(TEXTLITE_OUTPUT_TOKENS_LABEL, total_service_cost_dict, tenant_attribution_percentage_json) 11 | """ 12 | 13 | replace_code_in_file("../../cdk/lib/tenant-template/services/aggregate-metrics/invoke_model_tenant_cost.py", original_tenant_cost_service_str, update_tenant_cost_service_str) 14 | 15 | -------------------------------------------------------------------------------- /scripts/labs/lab4_get_total_service_cost.py: -------------------------------------------------------------------------------- 1 | from replace_code import replace_code_in_file 2 | 3 | original_str = """# TODO: Lab4 - Get total input and output tokens cost""" 4 | 5 | update_str = """# TODO: Lab4 - Get total input and output tokens cost 6 | if line_item in (EMBEDDING_TITAN_INPUT_TOKENS_LABEL, TEXTLITE_INPUT_TOKENS_LABEL,TEXTLITE_OUTPUT_TOKENS_LABEL): 7 | total_service_cost_dict[line_item] = cost 8 | """ 9 | 10 | replace_code_in_file("../../cdk/lib/tenant-template/services/aggregate-metrics/invoke_model_tenant_cost.py", original_str, update_str) 11 | -------------------------------------------------------------------------------- /scripts/labs/replace_code.py: -------------------------------------------------------------------------------- 1 | import re 2 | 3 | def replace_code_in_file(file_path, old_string, new_string): 4 | try: 5 | # Open the file for reading and writing 6 | with open(file_path, 'r') as file: 7 | content = file.read() 8 | 9 | # Replace the old string with the new string 10 | content = content.replace(old_string, new_string) 11 | 12 | # Write the modified content back to the file 13 | with open(file_path, 'w') as file: 14 | file.write(content) 15 | print(f"Successfully replaced code.") 16 | 17 | except FileNotFoundError: 18 | print(f"The file {file_path} was not found.") 19 | except Exception as e: 20 | print(f"An error occurred: {e}") 21 | 22 | def replace_code_in_file_regex(file_path, old_pattern, new_string): 23 | try: 24 | # Open the file for reading 25 | with open(file_path, 'r') as file: 26 | content = file.read() 27 | 28 | # Use regex to search for knowledge_base_id with any value 29 | content = re.sub(old_pattern, new_string, content, flags=re.DOTALL) 30 | 31 | # Write the modified content back to the file 32 | with open(file_path, 'w') as file: 33 | file.write(content) 34 | 35 | print("Successfully replaced code.") 36 | 37 | except FileNotFoundError: 38 | print(f"The file {file_path} was not found.") 39 | except Exception as e: 40 | print(f"An error occurred: {e}") -------------------------------------------------------------------------------- /scripts/provision-tenant.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash -e 2 | 3 | # Enable nocasematch option 4 | shopt -s nocasematch 5 | 6 | S3_TENANT_SOURCECODE_BUCKET_URL=$(aws cloudformation describe-stacks --stack-name saas-genai-workshop-bootstrap-template --query "Stacks[0].Outputs[?OutputKey=='TenantSourceCodeS3Bucket'].OutputValue" --output text) 7 | export CDK_PARAM_CODE_REPOSITORY_NAME="saas-genai-workshop" 8 | 9 | # Download the folder from S3 to local directory 10 | echo "Downloading folder from s3://$S3_TENANT_SOURCECODE_BUCKET_URL to $CDK_PARAM_CODE_REPOSITORY_NAME..." 11 | aws s3 cp "s3://$S3_TENANT_SOURCECODE_BUCKET_URL" "$CDK_PARAM_CODE_REPOSITORY_NAME" --recursive \ 12 | --exclude "cdk/cdk.out/*" --exclude "cdk/node_modules/*" --exclude ".git/*" --quiet 13 | cd $CDK_PARAM_CODE_REPOSITORY_NAME/cdk 14 | 15 | # Parse tenant details from the input message from step function 16 | export CDK_PARAM_TENANT_ID=$(echo $tenantId | tr -d '"') 17 | export CDK_PARAM_TENANT_NAME=$(echo $tenantName | tr -d '"') 18 | export TENANT_ADMIN_EMAIL=$(echo $email | tr -d '"') 19 | 20 | # Define variables 21 | STACK_NAME="saas-genai-workshop-bootstrap-template" 22 | USER_POOL_OUTPUT_PARAM_NAME="TenantUserpoolId" 23 | APP_CLIENT_ID_OUTPUT_PARAM_NAME="UserPoolClientId" 24 | API_GATEWAY_URL_OUTPUT_PARAM_NAME="ApiGatewayUrl" 25 | API_GATEWAY_USAGE_PLAN_ID_OUTPUT_PARAM_NAME="ApiGatewayUsagePlan" 26 | S3_PARAM_NAME="SaaSGenAIWorkshopS3Bucket" 27 | INGESTION_LAMBDA_ARN_PARAM_NAME="SaaSGenAIWorkshopTriggerIngestionLambdaArn" 28 | OSSC_ARN_PARAM_NAME="SaaSGenAIWorkshopOSSCollectionArn" 29 | INPUT_TOKENS="10000" 30 | OUTPUT_TOKENS="500" 31 | 32 | 33 | # Read tenant details from the cloudformation 34 | export REGION=$(aws configure get region) 35 | export SAAS_APP_USERPOOL_ID=$(aws cloudformation describe-stacks --stack-name $STACK_NAME --query "Stacks[0].Outputs[?OutputKey=='$USER_POOL_OUTPUT_PARAM_NAME'].OutputValue" --output text) 36 | export SAAS_APP_CLIENT_ID=$(aws cloudformation describe-stacks --stack-name $STACK_NAME --query "Stacks[0].Outputs[?OutputKey=='$APP_CLIENT_ID_OUTPUT_PARAM_NAME'].OutputValue" --output text) 37 | export API_GATEWAY_URL=$(aws cloudformation describe-stacks --stack-name $STACK_NAME --query "Stacks[0].Outputs[?OutputKey=='$API_GATEWAY_URL_OUTPUT_PARAM_NAME'].OutputValue" --output text) 38 | export API_GATEWAY_USAGE_PLAN_ID=$(aws cloudformation describe-stacks --stack-name $STACK_NAME --query "Stacks[0].Outputs[?OutputKey=='$API_GATEWAY_USAGE_PLAN_ID_OUTPUT_PARAM_NAME'].OutputValue" --output text) 39 | export S3_BUCKET=$(aws cloudformation describe-stacks --stack-name $STACK_NAME --query "Stacks[0].Outputs[?OutputKey=='$S3_PARAM_NAME'].OutputValue" --output text) 40 | export TRIGGER_PIPELINE_INGESTION_LAMBDA_ARN=$(aws cloudformation describe-stacks --stack-name $STACK_NAME --query "Stacks[0].Outputs[?OutputKey=='$INGESTION_LAMBDA_ARN_PARAM_NAME'].OutputValue" --output text) 41 | export OPENSEARCH_SERVERLESS_COLLECTION_ARN=$(aws cloudformation describe-stacks --stack-name $STACK_NAME --query "Stacks[0].Outputs[?OutputKey=='$OSSC_ARN_PARAM_NAME'].OutputValue" --output text) 42 | 43 | # Create Tenant API Key 44 | generate_api_key() { 45 | local suffix=${1:-sbt} 46 | local uuid=$(python3 -c "import uuid; print(uuid.uuid4())") 47 | echo "${uuid}-${suffix}" 48 | } 49 | 50 | TENANT_API_KEY=$(generate_api_key) 51 | 52 | # Error handling function 53 | check_error() { 54 | provision_script_name=$1 55 | exit_code=$2 56 | provision_output=$3 57 | if [[ "$exit_code" -ne 0 ]]; then 58 | echo "$provision_output" 59 | echo "ERROR: $provision_script_name failed. Exiting" 60 | exit 1 61 | fi 62 | echo "$provision_script_name completed successfully" 63 | } 64 | 65 | # Invoke tenant provisioning service 66 | pip3 install -r lib/tenant-template/tenant-provisioning/requirements.txt 67 | 68 | provision_name="Tenant Provisioning" 69 | # TODO: Lab1 - Add tenant provisioning service 70 | 71 | 72 | export KNOWLEDGE_BASE_NAME=$CDK_PARAM_TENANT_ID 73 | 74 | # List all knowledge bases and filter the results based on the KnowledgeBase name 75 | export KNOWLEDGE_BASE_ID=$(aws bedrock-agent list-knowledge-bases | jq -r '.[] | .[] | select(.name == $name) | .knowledgeBaseId' --arg name $KNOWLEDGE_BASE_NAME) 76 | 77 | # Create tenant admin user 78 | provision_name="Tenant Admin User Provisioning" 79 | # TODO: Lab1 - Uncomment below lines - user management service 80 | tenant_admin_output=$(python3 lib/tenant-template/user-management/user_management_service.py --tenant-id $CDK_PARAM_TENANT_ID --email $TENANT_ADMIN_EMAIL --user-role "TenantAdmin" 2>&1 >/dev/null && exit_code=$?) || exit_code=$? 81 | check_error "$provision_name" $exit_code "$tenant_admin_output" 82 | 83 | # Create JSON response of output parameters 84 | export tenantConfig=$(jq --arg SAAS_APP_USERPOOL_ID "$SAAS_APP_USERPOOL_ID" \ 85 | --arg SAAS_APP_CLIENT_ID "$SAAS_APP_CLIENT_ID" \ 86 | --arg API_GATEWAY_URL "$API_GATEWAY_URL" \ 87 | --arg TENANT_API_KEY "$TENANT_API_KEY" \ 88 | --arg CDK_PARAM_TENANT_NAME "$CDK_PARAM_TENANT_NAME" \ 89 | --arg KNOWLEDGE_BASE_ID "$KNOWLEDGE_BASE_ID" \ 90 | --arg INPUT_TOKENS "$INPUT_TOKENS" \ 91 | --arg OUTPUT_TOKENS "$OUTPUT_TOKENS" \ 92 | -n '{"tenantName":$CDK_PARAM_TENANT_NAME,"userPoolId":$SAAS_APP_USERPOOL_ID,"appClientId":$SAAS_APP_CLIENT_ID,"apiGatewayUrl":$API_GATEWAY_URL,"apiKey":$TENANT_API_KEY, "knowledgeBaseId":$KNOWLEDGE_BASE_ID, "inputTokens":$INPUT_TOKENS, "outputTokens":$OUTPUT_TOKENS}') 93 | export tenantStatus="Complete" 94 | --------------------------------------------------------------------------------