Skip to content

Commit

Permalink
WIP: revert Bedrock KB to agent only state
Browse files Browse the repository at this point in the history
  • Loading branch information
@san99tiago
san99tiago committed Oct 26, 2024
1 parent e729b85 commit 1d5ffa1
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 66 deletions.
129 changes: 65 additions & 64 deletions cdk/stacks/cdk_chatbot_api_stack.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -660,7 +660,7 @@ def create_bedrock_components(self) -> None:
name="encryption-policy",
policy='{"Rules":[{"ResourceType":"collection","Resource":["collection/*"]}],"AWSOwnedKey":true}',
type="encryption",
description="the encryption policy for the opensearch serverless collection",
description="Encryption policy for the opensearch serverless collection",
)

# We also need a security policy of type network so that the collection becomes accessable. The policy must be a string and the resource contains the collections it is applied to.
Expand All @@ -670,7 +670,7 @@ def create_bedrock_components(self) -> None:
name="network-policy",
policy='[{"Description":"Public access for collection","Rules":[{"ResourceType":"dashboard","Resource":["collection/*"]},{"ResourceType":"collection","Resource":["collection/*"]}],"AllowFromPublic":true}]',
type="network",
description="the network policy for the opensearch serverless collection",
description="Network policy for the opensearch serverless collection",
)

# Create the OpenSearch Collection
Expand Down Expand Up @@ -755,13 +755,14 @@ def create_bedrock_components(self) -> None:
self,
"OpenSearchServerlessAccessPolicy",
name=f"{self.main_resources_name}-data-access-policy",
policy=f'[{{"Description":"Access for bedrock","Rules":[{{"ResourceType":"index","Resource":["index/*/*"],"Permission":["aoss:*"]}},{{"ResourceType":"collection","Resource":["collection/*"],"Permission":["aoss:*"]}}],"Principal":["{bedrock_agent_role.role_arn}","{bedrock_kb_role.role_arn}","{create_index_lambda.role.role_arn}"]}}]',
policy=f'[{{"Description":"Access for bedrock","Rules":[{{"ResourceType":"index","Resource":["index/*/*"],"Permission":["aoss:*"]}},{{"ResourceType":"collection","Resource":["collection/*"],"Permission":["aoss:*"]}}],"Principal":["{bedrock_agent_role.role_arn}","{bedrock_kb_role.role_arn}","{create_index_lambda.role.role_arn}","arn:aws:iam::{self.account}:root"]}}]',
type="data",
description="the data access policy for the opensearch serverless collection",
description="Data access policy for the opensearch serverless collection",
)

opensearch_serverless_access_policy.add_dependency(
opensearch_serverless_collection
# Add dependencies to the collection
opensearch_serverless_collection.add_dependency(
opensearch_serverless_access_policy
)

# Define the request body for the lambda invoke api call that the custom resource will use
Expand Down Expand Up @@ -800,58 +801,58 @@ 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)
# # 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)
Expand Down Expand Up @@ -891,12 +892,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
4 changes: 2 additions & 2 deletions custom_resources/create_oss_index.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,8 @@ def handler(event, context):
"type": "knn_vector",
"dimension": 1536,
"method": {
"space_type": "innerproduct",
"engine": "FAISS",
"space_type": "l2",
"engine": "faiss",
"name": "hnsw",
"parameters": {
"m": 16,
Expand Down

0 comments on commit 1d5ffa1

Please sign in to comment.