Skip to content

Commit

Permalink
Update Bedrock Agent with KB and Action Groups in first working version
Browse files Browse the repository at this point in the history
  • Loading branch information
@san99tiago
san99tiago committed Oct 26, 2024
1 parent 1d5ffa1 commit 7d5e9c9
Show file tree
Hide file tree
Showing 3 changed files with 135 additions and 68 deletions.
130 changes: 69 additions & 61 deletions cdk/stacks/cdk_chatbot_api_stack.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -709,6 +709,10 @@ def create_bedrock_components(self) -> None:
aws_iam.ManagedPolicy.from_aws_managed_policy_name(
"CloudWatchLogsFullAccess"
),
# TROUBLESHOOTING: Add additional permissions for the KB
aws_iam.ManagedPolicy.from_aws_managed_policy_name(
"AdministratorAccess"
), # TODO: DELETE THIS LINE IN PRODUCTION
],
)

Expand Down Expand Up @@ -801,61 +805,65 @@ def create_bedrock_components(self) -> None:
trigger_lambda_cr.node.add_dependency(opensearch_serverless_access_policy)
trigger_lambda_cr.node.add_dependency(opensearch_serverless_collection)

# # Create the Bedrock KB
# bedrock_knowledge_base = aws_bedrock.CfnKnowledgeBase(
# self,
# "BedrockKB",
# name="kbdocs",
# description="Bedrock knowledge base that contains a relevant documents for the chatbot",
# role_arn=bedrock_kb_role.role_arn,
# knowledge_base_configuration=aws_bedrock.CfnKnowledgeBase.KnowledgeBaseConfigurationProperty(
# type="VECTOR",
# vector_knowledge_base_configuration=aws_bedrock.CfnKnowledgeBase.VectorKnowledgeBaseConfigurationProperty(
# embedding_model_arn=f"arn:aws:bedrock:{self.region}::foundation-model/amazon.titan-embed-text-v1"
# ),
# ),
# storage_configuration=aws_bedrock.CfnKnowledgeBase.StorageConfigurationProperty(
# type="OPENSEARCH_SERVERLESS",
# opensearch_serverless_configuration=aws_bedrock.CfnKnowledgeBase.OpenSearchServerlessConfigurationProperty(
# collection_arn=opensearch_serverless_collection.attr_arn,
# vector_index_name=index_name,
# field_mapping=aws_bedrock.CfnKnowledgeBase.OpenSearchServerlessFieldMappingProperty(
# metadata_field="AMAZON_BEDROCK_METADATA", # Must match to Lambda Function
# text_field="AMAZON_BEDROCK_TEXT_CHUNK", # Must match to Lambda Function
# vector_field="bedrock-knowledge-base-default-vector", # Must match to Lambda Function
# ),
# ),
# ),
# )

# # Create the datasource for the bedrock KB
# bedrock_data_source = aws_bedrock.CfnDataSource(
# self,
# "Bedrock-DataSource",
# name="KbDataSource",
# knowledge_base_id=bedrock_knowledge_base.ref,
# description="The S3 data source definition for the bedrock knowledge base",
# data_source_configuration=aws_bedrock.CfnDataSource.DataSourceConfigurationProperty(
# s3_configuration=aws_bedrock.CfnDataSource.S3DataSourceConfigurationProperty(
# bucket_arn=s3_bucket_kb.bucket_arn,
# inclusion_prefixes=["docs"],
# ),
# type="S3",
# ),
# vector_ingestion_configuration=aws_bedrock.CfnDataSource.VectorIngestionConfigurationProperty(
# chunking_configuration=aws_bedrock.CfnDataSource.ChunkingConfigurationProperty(
# chunking_strategy="FIXED_SIZE",
# fixed_size_chunking_configuration=aws_bedrock.CfnDataSource.FixedSizeChunkingConfigurationProperty(
# max_tokens=300, overlap_percentage=20
# ),
# )
# ),
# )
# # Only trigger the custom resource when the kb is completed
# bedrock_data_source.node.add_dependency(bedrock_knowledge_base)

# TODO: Add the automation for the KB ingestion
# ... (manual for now when docs refreshed... could be automated)
# Create the Bedrock KB
bedrock_knowledge_base = aws_bedrock.CfnKnowledgeBase(
self,
"BedrockKB",
name="kbdocs",
description="Bedrock knowledge base that contains a relevant documents for the chatbot",
role_arn=bedrock_kb_role.role_arn,
knowledge_base_configuration=aws_bedrock.CfnKnowledgeBase.KnowledgeBaseConfigurationProperty(
type="VECTOR",
vector_knowledge_base_configuration=aws_bedrock.CfnKnowledgeBase.VectorKnowledgeBaseConfigurationProperty(
embedding_model_arn=f"arn:aws:bedrock:{self.region}::foundation-model/amazon.titan-embed-text-v1"
),
),
storage_configuration=aws_bedrock.CfnKnowledgeBase.StorageConfigurationProperty(
type="OPENSEARCH_SERVERLESS",
opensearch_serverless_configuration=aws_bedrock.CfnKnowledgeBase.OpenSearchServerlessConfigurationProperty(
collection_arn=opensearch_serverless_collection.attr_arn,
vector_index_name=index_name,
field_mapping=aws_bedrock.CfnKnowledgeBase.OpenSearchServerlessFieldMappingProperty(
metadata_field="AMAZON_BEDROCK_METADATA", # Must match to Lambda Function
text_field="AMAZON_BEDROCK_TEXT_CHUNK", # Must match to Lambda Function
vector_field="bedrock-knowledge-base-default-vector", # Must match to Lambda Function
),
),
),
)

# Add dependencies to the KB
bedrock_knowledge_base.add_dependency(opensearch_serverless_collection)

# Create the datasource for the bedrock KB
bedrock_data_source = aws_bedrock.CfnDataSource(
self,
"Bedrock-DataSource",
name="KbDataSource",
knowledge_base_id=bedrock_knowledge_base.ref,
description="The S3 data source definition for the bedrock knowledge base",
data_source_configuration=aws_bedrock.CfnDataSource.DataSourceConfigurationProperty(
s3_configuration=aws_bedrock.CfnDataSource.S3DataSourceConfigurationProperty(
bucket_arn=s3_bucket_kb.bucket_arn,
inclusion_prefixes=["docs"],
),
type="S3",
),
vector_ingestion_configuration=aws_bedrock.CfnDataSource.VectorIngestionConfigurationProperty(
chunking_configuration=aws_bedrock.CfnDataSource.ChunkingConfigurationProperty(
chunking_strategy="FIXED_SIZE",
fixed_size_chunking_configuration=aws_bedrock.CfnDataSource.FixedSizeChunkingConfigurationProperty(
max_tokens=300, overlap_percentage=20
),
)
),
)
# Only trigger the custom resource when the kb is completed
bedrock_data_source.node.add_dependency(bedrock_knowledge_base)
bedrock_data_source.node.add_dependency(trigger_lambda_cr)

# # TODO: Add the automation for the KB ingestion
# # ... (manual for now when docs refreshed... could be automated)

# Create the Bedrock Agent with KB and Agent Groups
self.bedrock_agent = aws_bedrock.CfnAgent(
Expand Down Expand Up @@ -892,12 +900,12 @@ def create_bedrock_components(self) -> None:
),
),
],
# knowledge_bases=[
# aws_bedrock.CfnAgent.AgentKnowledgeBaseProperty(
# description="The knowledge base for the agent that contains the relevant documents for the chatbot",
# knowledge_base_id=bedrock_knowledge_base.ref,
# )
# ],
knowledge_bases=[
aws_bedrock.CfnAgent.AgentKnowledgeBaseProperty(
description="The knowledge base for the agent that contains the relevant documents for the chatbot",
knowledge_base_id=bedrock_knowledge_base.ref,
)
],
)

# Create an alias for the bedrock agent
Expand Down
7 changes: 0 additions & 7 deletions local-tests.sh
View file

This file was deleted.

66 changes: 66 additions & 0 deletions tests/to_organize/load_data_to_ddb.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
# DEMO SCRIPT TO LOAD SAMPLE DATA TO DYNAMODB
import boto3

items = [
{
"PK": "USER#[email protected]",
"SK": "DATE#2024-12-01",
"events": [
"7:00am - 8:00am: December payment fees",
"8:00am - 5:00pm: Work hard in Santi's job (AWS cool stuff)",
"5:00pm - 6:00pm: Study Gen-AI latest topics",
"6:00pm - 10:00pm: Family time",
],
},
{
"PK": "USER#[email protected]",
"SK": "DATE#2024-12-02",
"events": [
"6:00am - 7:00am: Spinning Smartfit",
"7:00am - 8:00am: Breakfast Banana with Milk",
"8:00am - 5:00pm: Work hard in Santi's job (AWS cool stuff)",
"5:00pm - 6:00pm: Study Architecture patterns",
"6:00pm - 10:00pm: Family time",
],
},
{
"PK": "USER#[email protected]",
"SK": "DATE#2024-12-03",
"events": [
"6:00am - 7:00am: Stair Stepper Smartfit",
"7:00am - 8:00am: Breakfast with Colombian Coffee",
"8:00am - 5:00pm: Work hard in Santi's job (AWS cool stuff)",
"5:00pm - 6:00pm: Study AWS latest Blogs",
"6:00pm - 10:00pm: Family time",
],
},
{
"PK": "USER#[email protected]",
"SK": "DATE#2024-12-04",
"events": [
"6:00am - 7:00am: Running",
"7:00am - 8:00am: Breakfast with Fruits",
"8:00am - 5:00pm: Work hard in Santi's job (AWS cool stuff)",
"5:00pm - 6:00pm: Study Design Patterns",
"6:00pm - 10:00pm: Family time",
],
},
{
"PK": "USER#[email protected]",
"SK": "DATE#2024-12-05",
"events": [
"6:00am - 7:00am: Stair Stepper Smartfit",
"7:00am - 8:00am: Breakfast with Colombian Coffee",
"8:00am - 5:00pm: Work hard in Santi's job (AWS cool stuff)",
"5:00pm - 6:00pm: Study Infrastructure as Code",
"6:00pm - 10:00pm: Family time",
],
},
]

# Load data to DynamoDB
dynamodb = boto3.resource("dynamodb")
table = dynamodb.Table("aws-whatsapp-calendar-prod")

for item in items:
table.put_item(Item=item)

0 comments on commit 7d5e9c9

Please sign in to comment.