Skip to content

Commit d6aae56

Browse files
networkandcodenetworkandcode
committed
bedrock kb
1 parent 2aa7321 commit d6aae56

16 files changed

Lines changed: 233 additions & 33 deletions

‎structured-kb-demo/.env.example‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,10 @@ AWS_ACCESS_KEY_ID=
33
AWS_REGION=
44
AWS_SECRET_ACCESS_KEY=
55

6+
BEDROCK_KB=
7+
BEDROCK_KB_IAM_POLICY=
8+
BEDROCK_KB_IAM_ROLE=
9+
610
GLUE_CRAWLER=
711
GLUE_CRAWLER_IAM_POLICY=
812
GLUE_CRAWLER_IAM_ROLE=
@@ -13,6 +17,6 @@ REDSHIFT_NAMESPACE=
1317
REDSHIFT_WORKGROUP=
1418

1519
S3_BUCKET=
16-
S3_BUCKET_FOLDER=
20+
S3_FOLDER=
1721

1822
SQS_QUEUE=

‎structured-kb-demo/arns.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,31 @@
1+
import boto3
2+
from botocore.exceptions import ClientError
3+
4+
from logger import logger
15
from vars import (
26
AWS_ACCOUNT_ID,
37
S3_BUCKET,
48
GLUE_CRAWLER_IAM_POLICY,
59
GLUE_CRAWLER_IAM_ROLE,
610
REDSHIFT_IAM_ROLE,
11+
REDSHIFT_WORKGROUP,
12+
BEDROCK_KB_IAM_POLICY,
13+
BEDROCK_KB_IAM_ROLE,
714
AWS_REGION,
815
SQS_QUEUE,
916
)
1017

1118
AWS_MANAGED_GLUE_IAM_POLICY_ARN = "arn:aws:iam::aws:policy/service-role/AWSGlueServiceRole"
1219
AWS_MANAGED_REDSHIFT_IAM_POLICY_ARN = "arn:aws:iam::aws:policy/AmazonRedshiftAllCommandsFullAccess"
20+
21+
BEDROCK_KB_IAM_POLICY_ARN = f"arn:aws:iam::{AWS_ACCOUNT_ID}:policy/{BEDROCK_KB_IAM_POLICY}"
22+
BEDROCK_KB_IAM_ROLE_ARN = f"arn:aws:iam::{AWS_ACCOUNT_ID}:role/{BEDROCK_KB_IAM_ROLE}"
23+
1324
S3_BUCKET_ARN = f"arn:aws:s3:::{S3_BUCKET}"
25+
1426
GLUE_CRAWLER_IAM_POLICY_ARN = f"arn:aws:iam::{AWS_ACCOUNT_ID}:policy/{GLUE_CRAWLER_IAM_POLICY}"
1527
GLUE_CRAWLER_IAM_ROLE_ARN = f"arn:aws:iam::{AWS_ACCOUNT_ID}:role/{GLUE_CRAWLER_IAM_ROLE}"
28+
1629
SQS_QUEUE_ARN = f"arn:aws:sqs:{AWS_REGION}:{AWS_ACCOUNT_ID}:{SQS_QUEUE}"
30+
1731
REDSHIFT_IAM_ROLE_ARN = f"arn:aws:iam::{AWS_ACCOUNT_ID}:role/{REDSHIFT_IAM_ROLE}"

‎structured-kb-demo/get_redshift_wg_arn.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
try:
2+
client = boto3.client("redshift-serverless", region_name=AWS_REGION)
3+
response = client.get_workgroup(workgroupName=REDSHIFT_WORKGROUP)
4+
REDSHIFT_WORKGROUP_ARN = response["workgroup"]["workgroupArn"]
5+
except ClientError as e:
6+
logger.error(f"Error fetching Redshift workgroup ARN: {e.response["Error"]["Message"]}")
7+
REDSHIFT_WORKGROUP_ARN = None

‎structured-kb-demo/run_glue_crawler.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,10 @@
33
import boto3
44

55
from logger import logger
6-
from vars import CRAWLER, REGION
6+
from vars import GLUE_CRAWLER, AWS_REGION
77

88
def run_glue_crawler(crawler_name):
9-
glue = boto3.client("glue", region_name=REGION)
9+
glue = boto3.client("glue", region_name=AWS_REGION)
1010

1111
try:
1212
glue.start_crawler(Name=crawler_name)
@@ -31,4 +31,4 @@ def run_glue_crawler(crawler_name):
3131
except Exception as e:
3232
logger.error(f"Error: {str(e)}")
3333

34-
run_glue_crawler(CRAWLER)
34+
run_glue_crawler(GLUE_CRAWLER)

‎structured-kb-demo/setup_bedrock_kb.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
import boto3
2+
3+
from arns import BEDROCK_KB_IAM_ROLE_ARN, REDSHIFT_WORKGROUP_ARN
4+
from logger import logger
5+
from vars import AWS_REGION, BEDROCK_KB, GLUE_DB, S3_FOLDER
6+
7+
GLUE_TABLE = S3_FOLDER
8+
9+
bedrock = boto3.client("bedrock-agent", region_name=AWS_REGION)
10+
11+
bedrock.create_knowledge_base(
12+
name=BEDROCK_KB,
13+
roleArn=BEDROCK_KB_IAM_ROLE_ARN,
14+
knowledgeBaseConfiguration={
15+
"type": "STRUCTURED",
16+
"sqlKnowledgeBaseConfiguration": {
17+
"type": "REDSHIFT",
18+
"redshiftConfiguration": {
19+
"queryEngineConfiguration": {
20+
"type": "SERVERLESS",
21+
"serverlessConfiguration": {
22+
"workgroupArn": REDSHIFT_WORKGROUP_ARN,
23+
"authConfiguration": {
24+
"type": "IAM"
25+
}
26+
}
27+
},
28+
"storageConfigurations": [
29+
{
30+
"type": "AWS_DATA_CATALOG",
31+
"awsDataCatalogConfiguration": {
32+
"tableNames": [f"{GLUE_DB}.{GLUE_TABLE}"]
33+
}
34+
}
35+
]
36+
}
37+
}
38+
}
39+
)
40+
41+
logger.info(f"KB Created: {BEDROCK_KB} with Redshift Serverless as data source.")

‎structured-kb-demo/setup_bedrock_kb_iam_policy.py‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
import boto3
2+
import json
3+
4+
from logger import logger
5+
from vars import AWS_ACCOUNT_ID, BEDROCK_KB_IAM_POLICY
6+
7+
iam = boto3.client("iam")
8+
9+
# Define the policy document
10+
policy_document = {
11+
"Version": "2012-10-17",
12+
"Statement": [
13+
{
14+
"Sid": "RedshiftDataAPIStatementPermissions",
15+
"Effect": "Allow",
16+
"Action": [
17+
"redshift-data:GetStatementResult",
18+
"redshift-data:DescribeStatement",
19+
"redshift-data:CancelStatement"
20+
],
21+
"Resource": ["*"],
22+
"Condition": {
23+
"StringEquals": {
24+
"redshift-data:statement-owner-iam-userid": "${aws:userid}"
25+
}
26+
}
27+
},
28+
{
29+
"Sid": "RedshiftDataAPIExecutePermissions",
30+
"Effect": "Allow",
31+
"Action": [
32+
"redshift-data:ExecuteStatement"
33+
],
34+
"Resource": [
35+
f"arn:aws:redshift-serverless:us-east-1:{AWS_ACCOUNT_ID}:workgroup/*"
36+
]
37+
},
38+
{
39+
"Sid": "RedshiftServerlessGetCredentials",
40+
"Effect": "Allow",
41+
"Action": "redshift-serverless:GetCredentials",
42+
"Resource": [
43+
f"arn:aws:redshift-serverless:us-east-1:{AWS_ACCOUNT_ID}:workgroup/*"
44+
]
45+
},
46+
{
47+
"Sid": "SqlWorkbenchAccess",
48+
"Effect": "Allow",
49+
"Action": [
50+
"sqlworkbench:GetSqlRecommendations",
51+
"sqlworkbench:PutSqlGenerationContext",
52+
"sqlworkbench:GetSqlGenerationContext",
53+
"sqlworkbench:DeleteSqlGenerationContext"
54+
],
55+
"Resource": "*"
56+
},
57+
{
58+
"Sid": "KbAccess",
59+
"Effect": "Allow",
60+
"Action": [
61+
"bedrock:GenerateQuery"
62+
],
63+
"Resource": "*"
64+
}
65+
]
66+
}
67+
68+
try:
69+
# Create the policy
70+
response = iam.create_policy(
71+
PolicyName=BEDROCK_KB_IAM_POLICY,
72+
PolicyDocument=json.dumps(policy_document),
73+
Description="Permissions for Bedrock Structured KB to query Redshift Serverless."
74+
)
75+
76+
logger.info(f"Successfully created policy!")
77+
78+
except iam.exceptions.EntityAlreadyExistsException:
79+
logger.info(f"Policy already exists.")
80+
except Exception as e:
81+
logger.error(f"An error occurred: {e}")

‎structured-kb-demo/setup_bedrock_kb_iam_role.py‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
import boto3
2+
import json
3+
4+
from arns import BEDROCK_KB_IAM_POLICY_ARN
5+
from logger import logger
6+
from vars import AWS_ACCOUNT_ID, AWS_REGION, BEDROCK_KB_IAM_ROLE
7+
8+
iam = boto3.client("iam", region_name=AWS_REGION)
9+
10+
trust_policy = {
11+
"Version": "2012-10-17",
12+
"Statement": [
13+
{
14+
"Sid": "TrustPolicyStatement",
15+
"Effect": "Allow",
16+
"Principal": {
17+
"Service": "bedrock.amazonaws.com"
18+
},
19+
"Action": "sts:AssumeRole",
20+
"Condition": {
21+
"StringEquals": {
22+
"aws:SourceAccount": AWS_ACCOUNT_ID
23+
},
24+
"ArnLike": {
25+
"aws:SourceArn": f"arn:aws:bedrock:{AWS_REGION}:{AWS_ACCOUNT_ID}:knowledge-base/*"
26+
}
27+
}
28+
}
29+
]
30+
}
31+
32+
try:
33+
iam.create_role(
34+
RoleName=BEDROCK_KB_IAM_ROLE,
35+
AssumeRolePolicyDocument=json.dumps(trust_policy),
36+
Description='IAM Role for Bedrock Knowledge Base to access Redshift Serverless'
37+
)
38+
logger.info(f"Created role: {BEDROCK_KB_IAM_ROLE}")
39+
40+
iam.attach_role_policy(
41+
RoleName=BEDROCK_KB_IAM_ROLE,
42+
PolicyArn=BEDROCK_KB_IAM_POLICY_ARN
43+
)
44+
logger.info(f"Attached IAM policy to BedrockKB IAM role.")
45+
46+
except iam.exceptions.EntityAlreadyExistsException:
47+
logger.warning(f"Role {BEDROCK_KB_IAM_ROLE} already exists.")
48+
except Exception as e:
49+
logger.error(f"Error: {e}")

‎structured-kb-demo/setup_glue_crawler.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,22 @@
11
import boto3
22

3-
from arns import GLUE_CRAWLER_IAM_ROLE_ARN, QUEUE_ARN
3+
from arns import GLUE_CRAWLER_IAM_ROLE_ARN, SQS_QUEUE_ARN
44
from logger import logger
5-
from vars import BUCKET, CRAWLER, DB, FOLDER
5+
from vars import AWS_REGION, GLUE_CRAWLER, GLUE_DB, S3_BUCKET, S3_FOLDER
66

7-
S3_PATH = f"s3://{BUCKET}/{FOLDER}"
7+
S3_PATH = f"s3://{S3_BUCKET}/{S3_FOLDER}"
88

9-
glue = boto3.client("glue", region_name="us-east-1")
9+
glue = boto3.client("glue", region_name=AWS_REGION)
1010

1111
CRAWLER_CONFIG = {
1212
"Role": GLUE_CRAWLER_IAM_ROLE_ARN,
13-
"DatabaseName": DB,
13+
"DatabaseName": GLUE_DB,
1414
"Description": "Crawler for inventory data triggered by SQS events",
1515
"Targets": {
1616
"S3Targets": [
1717
{
1818
"Path": S3_PATH,
19-
"EventQueueArn": QUEUE_ARN,
19+
"EventQueueArn": SQS_QUEUE_ARN,
2020
}
2121
]
2222
},
@@ -28,10 +28,10 @@
2828
}
2929

3030
try:
31-
glue.create_crawler(Name=CRAWLER, **CRAWLER_CONFIG)
31+
glue.create_crawler(Name=GLUE_CRAWLER, **CRAWLER_CONFIG)
3232
logger.info("Crawler created successfully.")
3333
except glue.exceptions.AlreadyExistsException:
34-
glue.update_crawler(Name=CRAWLER, **CRAWLER_CONFIG)
34+
glue.update_crawler(Name=GLUE_CRAWLER, **CRAWLER_CONFIG)
3535
logger.info("Crawler already exists. Updated configuration successfully.")
3636
except Exception as e:
3737
logger.error(f"Error creating crawler: {e}")

‎structured-kb-demo/setup_glue_crawler_iam_policy.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
import json
44

5-
from arns import BUCKET_ARN, GLUE_CRAWLER_IAM_POLICY_ARN, QUEUE_ARN
5+
from arns import S3_BUCKET_ARN, GLUE_CRAWLER_IAM_POLICY_ARN, SQS_QUEUE_ARN
66

77
from logger import logger
88
from vars import GLUE_CRAWLER_IAM_POLICY
@@ -19,7 +19,7 @@
1919
"s3:ListBucket"
2020
],
2121
"Resource": [
22-
f"{BUCKET_ARN}/*"
22+
f"{S3_BUCKET_ARN}/*"
2323
]
2424
},
2525
{
@@ -32,7 +32,7 @@
3232
"sqs:ReceiveMessage",
3333
"sqs:SetQueueAttributes"
3434
],
35-
"Resource": QUEUE_ARN
35+
"Resource": SQS_QUEUE_ARN
3636
},
3737
]
3838
}

‎structured-kb-demo/setup_glue_crawler_iam_role.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
from arns import AWS_MANAGED_GLUE_IAM_POLICY_ARN, GLUE_CRAWLER_IAM_POLICY_ARN
66
from logger import logger
7-
from vars import ACCOUNT_ID, GLUE_CRAWLER_IAM_ROLE
7+
from vars import AWS_ACCOUNT_ID, GLUE_CRAWLER_IAM_ROLE
88

99
iam = boto3.client("iam")
1010

@@ -19,7 +19,7 @@
1919
"Action": "sts:AssumeRole",
2020
"Condition": {
2121
"StringEquals": {
22-
"aws:SourceAccount": ACCOUNT_ID
22+
"aws:SourceAccount": AWS_ACCOUNT_ID
2323
}
2424
}
2525
}

0 commit comments

Comments
 (0)