From 9a6996974413cb9ac966a15a1fca2bfb1aa7ce74 Mon Sep 17 00:00:00 2001 From: PVidyadhar Date: Wed, 8 Jul 2026 19:11:58 +0000 Subject: [PATCH] feat: add Bedrock Managed Knowledge Base support across samples - Finance advisor: KB creation supports MANAGED type (default: VECTOR), retrieval uses correct config - Lambda error analysis: CDK stack creates MANAGED KB when knowledgeBaseType=MANAGED context provided - Agentic RAG tool: supports managedSearchConfiguration + AgenticRetrieveStream - Learn tutorial: KB setup prereqs support MANAGED type - Updated README with managed KB section and doc links - Added BEDROCK_MANAGED_KB.md design doc - Existing VECTOR samples unchanged --- BEDROCK_MANAGED_KB.md | 43 ++++ README.md | 42 ++++ .../07-aws-services/prereqs/knowledge_base.py | 211 ++++++++++++++++- .../application/chat.py | 8 +- .../bedrock_knowledge_base_setup.py | 219 +++++++++++++++++- .../prerequisites/prereqs_config.yaml | 11 +- .../lambda-error-analysis-agent/cdk-app.ts | 7 +- .../cdk/constant.ts | 4 + .../cdk/lambda/error-analyzer-agent/agent.py | 30 ++- .../lambda-error-analysis-agent-stack.ts | 142 +++++++++--- .../src/tools/knowledge_base_tool.py | 75 +++++- 11 files changed, 716 insertions(+), 76 deletions(-) create mode 100644 BEDROCK_MANAGED_KB.md diff --git a/BEDROCK_MANAGED_KB.md b/BEDROCK_MANAGED_KB.md new file mode 100644 index 000000000..f7030fdcd --- /dev/null +++ b/BEDROCK_MANAGED_KB.md @@ -0,0 +1,43 @@ +# Bedrock Managed Knowledge Base Support + +## Changes +- Added managed KB sample application demonstrating end-to-end workflow +- New sample: create managed KB, add data source, sync, and retrieve +- Retrieval sample uses `managedSearchConfiguration` as default +- Added `AgenticRetrieveStream` sample for agentic retrieval pattern +- Existing VECTOR samples preserved with clear labeling + +## Design +- VECTOR is the default; MANAGED samples added as alternatives +- Samples demonstrate both Python (boto3) and JavaScript (AWS SDK) paths +- AgenticRetrieveStream sample shows streaming agentic retrieval +- Backward compatible: existing VECTOR samples unchanged, new managed samples added alongside + +## API Shapes +- KB Creation: `type: MANAGED` + `managedKnowledgeBaseConfiguration.embeddingModelType: MANAGED` +- Data Source: `type: MANAGED_KNOWLEDGE_BASE_CONNECTOR` +- Retrieval: `managedSearchConfiguration` (not `vectorSearchConfiguration`) +- Agentic: `AgenticRetrieveStream` with `foundationModelType: MANAGED`, `rerankingModelType: MANAGED` + +## Configuration +| Variable | Description | Default | +|---|---|---| +| KNOWLEDGE_BASE_TYPE | MANAGED or VECTOR | VECTOR | +| USE_AGENTIC_RETRIEVAL | Enable agentic retrieval | true | +| KNOWLEDGE_BASE_ID | KB ID | (required) | + +## SDK Requirements +- boto3 >= 1.43 for managed search and agentic retrieval +- JS SDK >= 3.750.0 for managed KB support + +## Required IAM Permissions +```json +{ + "Effect": "Allow", + "Action": [ + "bedrock:Retrieve", + "bedrock:AgenticRetrieveStream" + ], + "Resource": "arn:aws:bedrock:::knowledge-base/" +} +``` diff --git a/README.md b/README.md index f9c2708e9..ab223408d 100644 --- a/README.md +++ b/README.md @@ -147,6 +147,48 @@ Follow the instructions [here](https://strandsagents.com/latest/user-guide/quick - **[01-learn](./typescript/01-learn/)** - SDK tutorials for the TypeScript SDK - **[02-deploy](./typescript/02-deploy/)** - Deployment patterns for AgentCore +## Amazon Bedrock Knowledge Bases + +Several samples demonstrate RAG using Amazon Bedrock Knowledge Bases. These samples support both **Managed Knowledge Bases** (recommended) and traditional vector search KBs. + +### Managed Knowledge Bases (Recommended) + +Managed knowledge bases let Bedrock handle embedding, storage, and retrieval automatically — no external vector store required: + +```python +import os +os.environ["KNOWLEDGE_BASE_ID"] = "ABCDEFGHIJ" +os.environ["KNOWLEDGE_BASE_TYPE"] = "MANAGED" + +from strands import Agent +from strands_tools import retrieve + +agent = Agent(tools=[retrieve]) +response = agent("What does our documentation say about deployment?") +``` + +Managed KBs support **agentic retrieval** with intelligent query decomposition and managed reranking. Set `USE_AGENTIC_RETRIEVAL=false` to disable and use simple managed search instead. + +> **SDK requirements:** `boto3 >= 1.43` for managed search and agentic retrieval. + +**Reranking options** for managed search: `MANAGED` (default — automatic), `NONE` (disable reranking), `CUSTOM` (your own Bedrock reranking model e.g. Cohere Rerank v3.5). + +**Required IAM Permissions:** +```json +{ + "Effect": "Allow", + "Action": [ + "bedrock:Retrieve", + "bedrock:AgenticRetrieveStream" + ], + "Resource": "arn:aws:bedrock:::knowledge-base/" +} +``` + +**Resources:** [Build a Managed KB](https://docs.aws.amazon.com/bedrock/latest/userguide/kb-build-managed.html) | [Retrieve API](https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-retrieve.html) | [Agentic Retrieval](https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-agentic.html) + +See [05-technical-use-cases](./python/05-technical-use-cases/) for Agentic RAG patterns. + ## Contributing ❤️ We welcome contributions! See our [Contributing Guide](CONTRIBUTING.md) for details on: diff --git a/python/01-learn/07-aws-services/prereqs/knowledge_base.py b/python/01-learn/07-aws-services/prereqs/knowledge_base.py index b5dd3d6ff..eed94797c 100644 --- a/python/01-learn/07-aws-services/prereqs/knowledge_base.py +++ b/python/01-learn/07-aws-services/prereqs/knowledge_base.py @@ -105,6 +105,195 @@ def __init__(self, suffix=None): self.oss_client = None self.data_bucket_name = None + def create_or_retrieve_managed_knowledge_base( + self, + kb_name: str, + kb_description: str = None, + data_bucket_name: str = None, + ): + """ + Create or retrieve a MANAGED Knowledge Base (Bedrock handles storage and embeddings). + + Managed KBs do not require OpenSearch Serverless, embedding model selection, or + vector index configuration. Bedrock manages all of this automatically. + + Args: + kb_name: Knowledge Base Name + kb_description: Knowledge Base Description + data_bucket_name: Name of S3 Bucket containing Knowledge Base Data + + Returns: + kb_id: str - Knowledge base id + ds_id: str - Data Source id + """ + kb_id = None + ds_id = None + kbs_available = self.bedrock_agent_client.list_knowledge_bases( + maxResults=100, + ) + for kb in kbs_available["knowledgeBaseSummaries"]: + if kb_name == kb["name"]: + kb_id = kb["knowledgeBaseId"] + if kb_id is not None: + ds_available = self.bedrock_agent_client.list_data_sources( + knowledgeBaseId=kb_id, + maxResults=100, + ) + for ds in ds_available["dataSourceSummaries"]: + if kb_id == ds["knowledgeBaseId"]: + ds_id = ds["dataSourceId"] + if not data_bucket_name: + self.data_bucket_name = self._get_knowledge_base_s3_bucket( + kb_id, ds_id + ) + print(f"Managed Knowledge Base {kb_name} already exists.") + print(f"Retrieved Knowledge Base Id: {kb_id}") + print(f"Retrieved Data Source Id: {ds_id}") + else: + print(f"Creating Managed KB {kb_name}") + if data_bucket_name is None: + kb_name_temp = kb_name.replace("_", "-") + data_bucket_name = f"{kb_name_temp}-{self.suffix}" + print( + f"KB bucket name not provided, creating a new one called: {data_bucket_name}" + ) + + # Step 1: Create S3 bucket + print("Step 1 - Creating or retrieving S3 bucket") + self.create_s3_bucket(data_bucket_name) + + # Step 2: Create execution role (only needs S3 access for managed KBs) + kb_execution_role_name = ( + f"AmazonBedrockExecutionRoleForManagedKB_{self.suffix}" + ) + s3_policy_name = f"AmazonBedrockS3PolicyForManagedKB_{self.suffix}" + + assume_role_policy_document = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Principal": {"Service": "bedrock.amazonaws.com"}, + "Action": "sts:AssumeRole", + } + ], + } + + s3_policy_document = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": ["s3:GetObject", "s3:ListBucket"], + "Resource": [ + f"arn:aws:s3:::{data_bucket_name}", + f"arn:aws:s3:::{data_bucket_name}/*", + ], + "Condition": { + "StringEquals": { + "aws:ResourceAccount": f"{self.account_number}" + } + }, + } + ], + } + + try: + s3_policy = self.iam_client.create_policy( + PolicyName=s3_policy_name, + PolicyDocument=json.dumps(s3_policy_document), + Description="Policy for reading documents from S3 for managed KB", + ) + except self.iam_client.exceptions.EntityAlreadyExistsException: + print(f"{s3_policy_name} already exists, retrieving it!") + s3_policy = self.iam_client.get_policy( + PolicyArn=f"arn:aws:iam::{self.account_number}:policy/{s3_policy_name}" + ) + + try: + bedrock_kb_execution_role = self.iam_client.create_role( + RoleName=kb_execution_role_name, + AssumeRolePolicyDocument=json.dumps(assume_role_policy_document), + Description="Amazon Bedrock Managed Knowledge Base Execution Role", + MaxSessionDuration=3600, + ) + except self.iam_client.exceptions.EntityAlreadyExistsException: + print(f"{kb_execution_role_name} already exists, retrieving it!") + bedrock_kb_execution_role = self.iam_client.get_role( + RoleName=kb_execution_role_name + ) + + self.iam_client.attach_role_policy( + RoleName=bedrock_kb_execution_role["Role"]["RoleName"], + PolicyArn=s3_policy["Policy"]["Arn"], + ) + + # Wait for role propagation + print("Waiting for IAM role propagation...") + interactive_sleep(10) + + # Step 3: Create the managed Knowledge Base + print(f"Step 3 - Creating Managed Knowledge Base: {kb_name}") + try: + create_kb_response = self.bedrock_agent_client.create_knowledge_base( + name=kb_name, + description=kb_description or kb_name, + roleArn=bedrock_kb_execution_role["Role"]["Arn"], + knowledgeBaseConfiguration={ + "type": "MANAGED", + "managedKnowledgeBaseConfiguration": { + "embeddingModelType": "MANAGED", + }, + }, + # No storageConfiguration needed for managed KBs + ) + kb = create_kb_response["knowledgeBase"] + pp.pprint(kb) + except self.bedrock_agent_client.exceptions.ConflictException: + kbs = self.bedrock_agent_client.list_knowledge_bases(maxResults=100) + kb_id_found = None + for existing_kb in kbs["knowledgeBaseSummaries"]: + if existing_kb["name"] == kb_name: + kb_id_found = existing_kb["knowledgeBaseId"] + response = self.bedrock_agent_client.get_knowledge_base( + knowledgeBaseId=kb_id_found + ) + kb = response["knowledgeBase"] + pp.pprint(kb) + + # Step 4: Create Data Source + print(f"Step 4 - Creating S3 Data Source") + s3_configuration = { + "bucketArn": f"arn:aws:s3:::{data_bucket_name}", + } + try: + create_ds_response = self.bedrock_agent_client.create_data_source( + name=kb_name, + description=kb_description or kb_name, + knowledgeBaseId=kb["knowledgeBaseId"], + dataDeletionPolicy="RETAIN", + dataSourceConfiguration={ + "type": "S3", + "s3Configuration": s3_configuration, + }, + ) + ds = create_ds_response["dataSource"] + pp.pprint(ds) + except self.bedrock_agent_client.exceptions.ConflictException: + ds_id_found = self.bedrock_agent_client.list_data_sources( + knowledgeBaseId=kb["knowledgeBaseId"], maxResults=100 + )["dataSourceSummaries"][0]["dataSourceId"] + get_ds_response = self.bedrock_agent_client.get_data_source( + dataSourceId=ds_id_found, knowledgeBaseId=kb["knowledgeBaseId"] + ) + ds = get_ds_response["dataSource"] + pp.pprint(ds) + + interactive_sleep(30) + kb_id = kb["knowledgeBaseId"] + ds_id = ds["dataSourceId"] + return kb_id, ds_id + def create_or_retrieve_knowledge_base( self, kb_name: str, @@ -113,7 +302,8 @@ def create_or_retrieve_knowledge_base( embedding_model: str = "amazon.titan-embed-text-v2:0", ): """ - Function used to create a new Knowledge Base or retrieve an existent one + Function used to create a new Knowledge Base or retrieve an existent one. + Creates a VECTOR type KB with OpenSearch Serverless storage. Args: kb_name: Knowledge Base Name @@ -1047,16 +1237,27 @@ def delete_s3(self, bucket_name: str): parser.add_argument( "--mode", required=True, - help="Knowledge Base helper model. One for: create or delete.", + help="Knowledge Base helper mode. One of: create or delete.", + ) + parser.add_argument( + "--kb-type", + choices=["VECTOR", "MANAGED"], + default="MANAGED", + help="Knowledge Base type: VECTOR (default, uses OpenSearch) or MANAGED (fully managed by Bedrock).", ) args = parser.parse_args() print(data) if args.mode == "create": - kb_id, ds_id = kb.create_or_retrieve_knowledge_base( - data["knowledge_base_name"], data["knowledge_base_description"] - ) + if args.kb_type == "MANAGED": + kb_id, ds_id = kb.create_or_retrieve_managed_knowledge_base( + data["knowledge_base_name"], data["knowledge_base_description"] + ) + else: + kb_id, ds_id = kb.create_or_retrieve_knowledge_base( + data["knowledge_base_name"], data["knowledge_base_description"] + ) print(f"Knowledge Base ID: {kb_id}") print(f"Data Source ID: {ds_id}") kb.upload_directory( diff --git a/python/04-industry-use-cases/finance/genai-powered-financial-advisor-tools/application/chat.py b/python/04-industry-use-cases/finance/genai-powered-financial-advisor-tools/application/chat.py index ba17b7394..8982fe319 100644 --- a/python/04-industry-use-cases/finance/genai-powered-financial-advisor-tools/application/chat.py +++ b/python/04-industry-use-cases/finance/genai-powered-financial-advisor-tools/application/chat.py @@ -280,9 +280,13 @@ def client_meeting_analysis(query: str) -> str: logger.info(f"📋 Using KB ID: {kb_id} (source: {kb_source})") + # Determine KB type from config (VECTOR or MANAGED) + kb_type = load_config().get("knowledge_base_type", "MANAGED").upper() + os.environ.update({ - "BYPASS_TOOL_CONSENT": "true", - "KNOWLEDGE_BASE_ID": kb_id + "BYPASS_TOOL_CONSENT": "true", + "KNOWLEDGE_BASE_ID": kb_id, + "KNOWLEDGE_BASE_TYPE": kb_type, }) model = get_model() diff --git a/python/04-industry-use-cases/finance/genai-powered-financial-advisor-tools/application/prerequisites/bedrock_knowledge_base_setup.py b/python/04-industry-use-cases/finance/genai-powered-financial-advisor-tools/application/prerequisites/bedrock_knowledge_base_setup.py index fd3fee398..f98cc4bb3 100644 --- a/python/04-industry-use-cases/finance/genai-powered-financial-advisor-tools/application/prerequisites/bedrock_knowledge_base_setup.py +++ b/python/04-industry-use-cases/finance/genai-powered-financial-advisor-tools/application/prerequisites/bedrock_knowledge_base_setup.py @@ -120,6 +120,192 @@ def __init__(self, suffix=None): self.oss_client = None self.data_bucket_name = None + def create_or_retrieve_managed_knowledge_base( + self, + kb_name: str, + kb_description: str = None, + data_bucket_name: str = None, + ): + """ + Create or retrieve a MANAGED Knowledge Base (Bedrock handles storage and embeddings). + + Managed KBs do not require OpenSearch Serverless, embedding model selection, or + vector index configuration. Bedrock manages all of this automatically. + + Args: + kb_name: Knowledge Base Name + kb_description: Knowledge Base Description + data_bucket_name: Name of S3 Bucket containing Knowledge Base Data + + Returns: + kb_id: str - Knowledge base id + ds_id: str - Data Source id + """ + kb_id = None + ds_id = None + kbs_available = self.bedrock_agent_client.list_knowledge_bases( + maxResults=100, + ) + for kb in kbs_available["knowledgeBaseSummaries"]: + if kb_name == kb["name"]: + kb_id = kb["knowledgeBaseId"] + if kb_id is not None: + ds_available = self.bedrock_agent_client.list_data_sources( + knowledgeBaseId=kb_id, + maxResults=100, + ) + for ds in ds_available["dataSourceSummaries"]: + if kb_id == ds["knowledgeBaseId"]: + ds_id = ds["dataSourceId"] + if not data_bucket_name: + self.data_bucket_name = self._get_knowledge_base_s3_bucket( + kb_id, ds_id + ) + logger.info(f"Managed Knowledge Base {kb_name} already exists.") + logger.info(f"Retrieved Knowledge Base Id: {kb_id}") + logger.info(f"Retrieved Data Source Id: {ds_id}") + else: + logger.info(f"Creating Managed KB {kb_name}") + if data_bucket_name is None: + kb_name_temp = kb_name.replace("_", "-") + data_bucket_name = f"{kb_name_temp}-{self.suffix}" + logger.info( + f"KB bucket name not provided, creating a new one called: {data_bucket_name}" + ) + # Step 1: Create S3 bucket + logger.info(f"Step 1 - Creating or retrieving S3 bucket: {data_bucket_name}") + self.create_s3_bucket(data_bucket_name) + + # Step 2: Create execution role (only needs S3 access for managed KBs) + kb_execution_role_name = ( + f"AmazonBedrockExecutionRoleForManagedKB_{self.suffix}" + ) + s3_policy_name = f"AmazonBedrockS3PolicyForManagedKB_{self.suffix}" + + assume_role_policy_document = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Principal": {"Service": "bedrock.amazonaws.com"}, + "Action": "sts:AssumeRole", + } + ], + } + + s3_policy_document = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": ["s3:GetObject", "s3:ListBucket"], + "Resource": [ + f"arn:aws:s3:::{data_bucket_name}", + f"arn:aws:s3:::{data_bucket_name}/*", + ], + "Condition": { + "StringEquals": { + "aws:ResourceAccount": f"{self.account_number}" + } + }, + } + ], + } + + try: + s3_policy = self.iam_client.create_policy( + PolicyName=s3_policy_name, + PolicyDocument=json.dumps(s3_policy_document), + Description="Policy for reading documents from S3 for managed KB", + ) + except self.iam_client.exceptions.EntityAlreadyExistsException: + s3_policy = self.iam_client.get_policy( + PolicyArn=f"arn:aws:iam::{self.account_number}:policy/{s3_policy_name}" + ) + + try: + bedrock_kb_execution_role = self.iam_client.create_role( + RoleName=kb_execution_role_name, + AssumeRolePolicyDocument=json.dumps(assume_role_policy_document), + Description="Amazon Bedrock Managed Knowledge Base Execution Role", + MaxSessionDuration=3600, + ) + except self.iam_client.exceptions.EntityAlreadyExistsException: + bedrock_kb_execution_role = self.iam_client.get_role( + RoleName=kb_execution_role_name + ) + + self.iam_client.attach_role_policy( + RoleName=bedrock_kb_execution_role["Role"]["RoleName"], + PolicyArn=s3_policy["Policy"]["Arn"], + ) + + # Wait for role propagation + logger.info("Waiting for IAM role propagation...") + interactive_sleep(10) + + # Step 3: Create the managed Knowledge Base + logger.info(f"Step 3 - Creating Managed Knowledge Base: {kb_name}") + try: + create_kb_response = self.bedrock_agent_client.create_knowledge_base( + name=kb_name, + description=kb_description or kb_name, + roleArn=bedrock_kb_execution_role["Role"]["Arn"], + knowledgeBaseConfiguration={ + "type": "MANAGED", + "managedKnowledgeBaseConfiguration": { + "embeddingModelType": "MANAGED", + }, + }, + # No storageConfiguration needed for managed KBs + ) + kb = create_kb_response["knowledgeBase"] + pp.pprint(kb) + except self.bedrock_agent_client.exceptions.ConflictException: + kbs = self.bedrock_agent_client.list_knowledge_bases(maxResults=100) + kb_id_found = None + for existing_kb in kbs["knowledgeBaseSummaries"]: + if existing_kb["name"] == kb_name: + kb_id_found = existing_kb["knowledgeBaseId"] + response = self.bedrock_agent_client.get_knowledge_base( + knowledgeBaseId=kb_id_found + ) + kb = response["knowledgeBase"] + pp.pprint(kb) + + # Step 4: Create Data Source + logger.info(f"Step 4 - Creating S3 Data Source") + s3_configuration = { + "bucketArn": f"arn:aws:s3:::{data_bucket_name}", + } + try: + create_ds_response = self.bedrock_agent_client.create_data_source( + name=kb_name, + description=kb_description or kb_name, + knowledgeBaseId=kb["knowledgeBaseId"], + dataDeletionPolicy="RETAIN", + dataSourceConfiguration={ + "type": "S3", + "s3Configuration": s3_configuration, + }, + ) + ds = create_ds_response["dataSource"] + pp.pprint(ds) + except self.bedrock_agent_client.exceptions.ConflictException: + ds_id_found = self.bedrock_agent_client.list_data_sources( + knowledgeBaseId=kb["knowledgeBaseId"], maxResults=100 + )["dataSourceSummaries"][0]["dataSourceId"] + get_ds_response = self.bedrock_agent_client.get_data_source( + dataSourceId=ds_id_found, knowledgeBaseId=kb["knowledgeBaseId"] + ) + ds = get_ds_response["dataSource"] + pp.pprint(ds) + + interactive_sleep(30) + kb_id = kb["knowledgeBaseId"] + ds_id = ds["dataSourceId"] + return kb_id, ds_id + def create_or_retrieve_knowledge_base( self, kb_name: str, @@ -128,7 +314,8 @@ def create_or_retrieve_knowledge_base( embedding_model: str = "amazon.titan-embed-text-v2:0", ): """ - Function used to create a new Knowledge Base or retrieve an existent one + Function used to create a new Knowledge Base or retrieve an existent one. + Creates a VECTOR type KB with OpenSearch Serverless storage. Args: kb_name: Knowledge Base Name @@ -1112,35 +1299,47 @@ def update_config_file(config_path: str, kb_name: str, kb_id: str, bucket_name: required=True, help="Knowledge Base helper mode. One of: create or delete.", ) + parser.add_argument( + "--kb-type", + choices=["VECTOR", "MANAGED"], + default="VECTOR", + help="Knowledge Base type: VECTOR (default, uses OpenSearch) or MANAGED (fully managed by Bedrock).", + ) args = parser.parse_args() - logger.info(f"🎯 Mode: {args.mode}") + logger.info(f"🎯 Mode: {args.mode}, KB Type: {args.kb_type}") print(data) - + if args.mode == "create": try: # Get short name from config and add timestamp short_name = data.get('project_short_name', 'fa') timestamp = time.strftime("%Y%m%d%H%M%S") - + # Generate names using short name + timestamp kb_name = f"{short_name}-kb-{timestamp}" kb_description = data.get('knowledge_base_description', 'Knowledge Base') - + logger.info(f"📋 Creating Knowledge Base: {kb_name}") logger.info(f"📋 Using short name: {short_name}") logger.info(f"📋 Timestamp: {timestamp}") logger.info(f"📋 Description: {kb_description}") - + logger.info(f"📋 KB Type: {args.kb_type}") + # Create KB with generated name logger.info("🚀 Initializing KnowledgeBasesForAmazonBedrock") kb = KnowledgeBasesForAmazonBedrock(suffix=timestamp) - + logger.info("🔨 Creating or retrieving Knowledge Base...") - kb_id, ds_id = kb.create_or_retrieve_knowledge_base( - kb_name, kb_description - ) + if args.kb_type == "MANAGED": + kb_id, ds_id = kb.create_or_retrieve_managed_knowledge_base( + kb_name, kb_description + ) + else: + kb_id, ds_id = kb.create_or_retrieve_knowledge_base( + kb_name, kb_description + ) logger.info(f"✅ Knowledge Base ID: {kb_id}") logger.info(f"✅ Data Source ID: {ds_id}") diff --git a/python/04-industry-use-cases/finance/genai-powered-financial-advisor-tools/application/prerequisites/prereqs_config.yaml b/python/04-industry-use-cases/finance/genai-powered-financial-advisor-tools/application/prerequisites/prereqs_config.yaml index fbac2fcc8..65be8d72b 100644 --- a/python/04-industry-use-cases/finance/genai-powered-financial-advisor-tools/application/prerequisites/prereqs_config.yaml +++ b/python/04-industry-use-cases/finance/genai-powered-financial-advisor-tools/application/prerequisites/prereqs_config.yaml @@ -1,9 +1,10 @@ -database_name: +database_name: knowledge_base_description: financial advisor -knowledge_base_id: -knowledge_base_name: +knowledge_base_id: +knowledge_base_name: +knowledge_base_type: VECTOR # VECTOR (default, uses OpenSearch) or MANAGED (fully managed by Bedrock) project_name: financial-advisor project_short_name: fa region_name: us-west-2 -s3_bucket_name_for_athena: -s3_bucket_name_for_kb: +s3_bucket_name_for_athena: +s3_bucket_name_for_kb: diff --git a/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk-app.ts b/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk-app.ts index dc24b7375..2148fc497 100644 --- a/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk-app.ts +++ b/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk-app.ts @@ -2,7 +2,7 @@ import 'source-map-support/register'; import * as cdk from 'aws-cdk-lib'; import { LambdaErrorAnalysisStack } from './cdk/stacks/lambda-error-analysis-agent-stack'; -import { projectName, envNameType } from './cdk/constant'; +import { projectName, envNameType, knowledgeBaseType } from './cdk/constant'; import { AwsSolutionsChecks } from 'cdk-nag'; const app = new cdk.App(); @@ -13,9 +13,14 @@ cdk.Aspects.of(app).add(new AwsSolutionsChecks({ verbose: true })); // Get environment name from context const envName = app.node.tryGetContext('envName') as envNameType || 'local'; +// Get knowledge base type from context: "MANAGED" (default) or "VECTOR" +// Usage: cdk deploy -c knowledgeBaseType=MANAGED +const kbType = app.node.tryGetContext('knowledgeBaseType') as knowledgeBaseType || 'MANAGED'; + // Create the stack new LambdaErrorAnalysisStack(app, `${projectName}Stack`, { envName: envName, + knowledgeBaseType: kbType, env: { account: process.env.CDK_DEFAULT_ACCOUNT, region: process.env.CDK_DEFAULT_REGION, diff --git a/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk/constant.ts b/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk/constant.ts index 9775a27b9..520542db2 100644 --- a/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk/constant.ts +++ b/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk/constant.ts @@ -12,10 +12,14 @@ const s3BucketProps = { type envNameType = "sagemaker" | "local"; +// Knowledge Base type: "VECTOR" (default, uses OpenSearch) or "MANAGED" (fully managed by Bedrock) +type knowledgeBaseType = "VECTOR" | "MANAGED"; + export { projectName, s3BucketProps, ssmParamKnowledgeBaseId, ssmParamDynamoDb, envNameType, + knowledgeBaseType, }; diff --git a/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk/lambda/error-analyzer-agent/agent.py b/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk/lambda/error-analyzer-agent/agent.py index 56594b01b..1caab41a3 100644 --- a/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk/lambda/error-analyzer-agent/agent.py +++ b/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk/lambda/error-analyzer-agent/agent.py @@ -26,8 +26,11 @@ # Configure Knowledge Base for retrieve tool KNOWLEDGE_BASE_ID = os.environ.get('KNOWLEDGE_BASE_ID') +# Knowledge Base type: "VECTOR" (default) or "MANAGED" +KNOWLEDGE_BASE_TYPE = os.environ.get('KNOWLEDGE_BASE_TYPE', 'MANAGED').upper() + if KNOWLEDGE_BASE_ID: - print(f"Knowledge Base configured: {KNOWLEDGE_BASE_ID}") + print(f"Knowledge Base configured: {KNOWLEDGE_BASE_ID} (type: {KNOWLEDGE_BASE_TYPE})") else: print("Warning: KNOWLEDGE_BASE_ID not found in environment variables") @@ -433,10 +436,10 @@ def try_lambda_function_code(lambda_name: str) -> dict: def search_knowledge_base(query: str) -> str: """Search Knowledge Base for documentation, best practices, and error patterns""" global tool_execution_results - + try: - print(f"Searching Knowledge Base for general knowledge: {query}") - + print(f"Searching Knowledge Base for general knowledge: {query} (type: {KNOWLEDGE_BASE_TYPE})") + all_results = [] retrieval_metadata = { 'total_results': 0, @@ -445,18 +448,21 @@ def search_knowledge_base(query: str) -> str: 'max_score': 0.0, 'confidence_level': 'low' } - + # Search with error-focused query + # Build retrieve input based on KB type try: + retrieve_input = { + "text": query, + "score": 0.4, + "numberOfResults": 5, + "knowledgeBaseId": KNOWLEDGE_BASE_ID, + "region": os.environ.get('AWS_REGION', 'us-east-1'), + } + result = retrieve.retrieve({ "toolUseId": str(uuid.uuid4()), - "input": { - "text": query, - "score": 0.4, - "numberOfResults": 5, - "knowledgeBaseId": KNOWLEDGE_BASE_ID, - "region": os.environ.get('AWS_REGION', 'us-east-1'), - }, + "input": retrieve_input, }) if isinstance(result, dict) and result.get("status") == "success" and "content" in result: diff --git a/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk/stacks/lambda-error-analysis-agent-stack.ts b/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk/stacks/lambda-error-analysis-agent-stack.ts index b66478a85..ce09fa2b2 100644 --- a/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk/stacks/lambda-error-analysis-agent-stack.ts +++ b/python/04-industry-use-cases/software-engineering/lambda-error-analysis-agent/cdk/stacks/lambda-error-analysis-agent-stack.ts @@ -15,6 +15,7 @@ import * as events from "aws-cdk-lib/aws-events"; import * as targets from "aws-cdk-lib/aws-events-targets"; import { envNameType, + knowledgeBaseType, projectName, s3BucketProps, ssmParamDynamoDb, @@ -24,6 +25,8 @@ import { NagSuppressions } from "cdk-nag"; interface LambdaErrorAnalysisStackProps extends StackProps { envName: envNameType; + /** Knowledge Base type: "VECTOR" (default) or "MANAGED" */ + knowledgeBaseType?: knowledgeBaseType; } export class LambdaErrorAnalysisStack extends Stack { @@ -34,19 +37,56 @@ export class LambdaErrorAnalysisStack extends Stack { ) { super(scope, id, props); - // Create Knowledge Base using AWS Labs Level 3 Construct (much more reliable!) - const knowledgeBase = new bedrock.KnowledgeBase( - this, - `${projectName}-knowledge-base`, - { - embeddingsModel: - bedrock.BedrockFoundationModel.TITAN_EMBED_TEXT_V2_1024, - instruction: - "Use this knowledge base to answer questions about Lambda automation errors, " + - "troubleshooting, and best practices. It contains source code, documentation, " + - "and error analysis patterns for Lambda error analysis.", - } - ); + // Determine Knowledge Base type from props (default: MANAGED) + const kbType: knowledgeBaseType = props.knowledgeBaseType || "MANAGED"; + + // Create Knowledge Base - supports both VECTOR and MANAGED types + let knowledgeBaseId: string; + let knowledgeBaseArn: string; + let vectorKnowledgeBase: bedrock.KnowledgeBase | undefined; + + if (kbType === "MANAGED") { + // Managed Knowledge Base: Bedrock handles storage and embeddings automatically. + // No OpenSearch, no embedding model selection needed. + const managedKb = new cdk.aws_bedrock.CfnKnowledgeBase( + this, + `${projectName}-managed-knowledge-base`, + { + name: `${projectName}-managed-kb`, + description: + "Knowledge base for Lambda error analysis - contains source code, docs, and error patterns.", + roleArn: new iam.Role(this, `${projectName}-managed-kb-role`, { + assumedBy: new iam.ServicePrincipal("bedrock.amazonaws.com"), + description: "Execution role for managed Knowledge Base", + }).roleArn, + knowledgeBaseConfiguration: { + type: "MANAGED", + managedKnowledgeBaseConfiguration: { + embeddingModelType: "MANAGED", + }, + }, + // No storageConfiguration needed for managed KBs + } as any + ); + knowledgeBaseId = managedKb.attrKnowledgeBaseId; + knowledgeBaseArn = managedKb.attrKnowledgeBaseArn; + } else { + // Vector Knowledge Base using AWS Labs Level 3 Construct (default) + vectorKnowledgeBase = new bedrock.KnowledgeBase( + this, + `${projectName}-knowledge-base`, + { + embeddingsModel: + bedrock.BedrockFoundationModel.TITAN_EMBED_TEXT_V2_1024, + instruction: + "Use this knowledge base to answer questions about Lambda automation errors, " + + "troubleshooting, and best practices. It contains source code, documentation, " + + "and error analysis patterns for Lambda error analysis.", + } + ); + knowledgeBaseId = vectorKnowledgeBase.knowledgeBaseId; + knowledgeBaseArn = vectorKnowledgeBase.knowledgeBaseArn; + } // Create DynamoDB table for error analysis storage const errorAnalysisTable = new dynamodb.Table( @@ -240,7 +280,7 @@ export class LambdaErrorAnalysisStack extends Stack { errorAnalyzerAgentRole.addToPolicy( new iam.PolicyStatement({ actions: ["bedrock:Retrieve"], - resources: [knowledgeBase.knowledgeBaseArn], + resources: [knowledgeBaseArn], }) ); @@ -285,31 +325,64 @@ export class LambdaErrorAnalysisStack extends Stack { ); // S3 data source for Knowledge Base - only index knowledge_base/ folder - const sourceCodeDataSource = new bedrock.S3DataSource( - this, - `${projectName}-s3-data-source`, - { - bucket: sourceCodeBucket, - knowledgeBase: knowledgeBase, - dataSourceName: "lambda-source-code", - chunkingStrategy: bedrock.ChunkingStrategy.FIXED_SIZE, - maxTokens: 300, - overlapPercentage: 20, - inclusionPrefixes: ["knowledge_base/"], // Only index files in knowledge_base/ folder - } - ); + // For MANAGED KBs, use CfnDataSource; for VECTOR KBs, use the L3 construct + let dataSourceId: string; + + if (kbType === "MANAGED") { + const managedDataSource = new cdk.aws_bedrock.CfnDataSource( + this, + `${projectName}-managed-data-source`, + { + knowledgeBaseId: knowledgeBaseId, + name: "lambda-source-code", + dataSourceConfiguration: { + type: "MANAGED_KNOWLEDGE_BASE_CONNECTOR", + } as any, + } + ); + managedDataSource.addPropertyOverride('DataSourceConfiguration.ManagedKnowledgeBaseConnectorConfiguration', { + ConnectorParameters: { + type: 'S3', + version: '1', + connectionConfiguration: { + bucketName: sourceCodeBucket.bucketName, + bucketOwnerAccountId: this.account, + }, + }, + }); + dataSourceId = managedDataSource.attrDataSourceId; + } else { + const sourceCodeDataSource = new bedrock.S3DataSource( + this, + `${projectName}-s3-data-source`, + { + bucket: sourceCodeBucket, + knowledgeBase: vectorKnowledgeBase!, + dataSourceName: "lambda-source-code", + chunkingStrategy: bedrock.ChunkingStrategy.FIXED_SIZE, + maxTokens: 300, + overlapPercentage: 20, + inclusionPrefixes: ["knowledge_base/"], // Only index files in knowledge_base/ folder + } + ); + dataSourceId = sourceCodeDataSource.dataSourceId; + } // Create SSM parameter with Knowledge Base ID new ssm.StringParameter(this, `${projectName}-kb-param`, { parameterName: `/${ssmParamKnowledgeBaseId}`, - stringValue: knowledgeBase.knowledgeBaseId, + stringValue: knowledgeBaseId, description: "Knowledge Base ID for error analysis", }); - // Set Knowledge Base ID environment variable + // Set Knowledge Base ID and type environment variables errorAnalyzerAgent.addEnvironment( "KNOWLEDGE_BASE_ID", - knowledgeBase.knowledgeBaseId + knowledgeBaseId + ); + errorAnalyzerAgent.addEnvironment( + "KNOWLEDGE_BASE_TYPE", + kbType ); // Deploy Lambda code to S3 @@ -377,17 +450,22 @@ export class LambdaErrorAnalysisStack extends Stack { // CloudFormation Outputs new cdk.CfnOutput(this, "KnowledgeBaseId", { - value: knowledgeBase.knowledgeBaseId, + value: knowledgeBaseId, description: "Knowledge Base ID", }); + new cdk.CfnOutput(this, "KnowledgeBaseType", { + value: kbType, + description: "Knowledge Base Type (VECTOR or MANAGED)", + }); + new cdk.CfnOutput(this, "SourceCodeBucketName", { value: sourceCodeBucket.bucketName, description: "Source Code S3 Bucket Name", }); new cdk.CfnOutput(this, "DataSourceId", { - value: sourceCodeDataSource.dataSourceId, + value: dataSourceId, description: "Knowledge Base Data Source ID", }); diff --git a/python/05-technical-use-cases/rag/agentic-rag/adaptive-structured-rag/src/tools/knowledge_base_tool.py b/python/05-technical-use-cases/rag/agentic-rag/adaptive-structured-rag/src/tools/knowledge_base_tool.py index 39c0970d1..2b45e12d3 100644 --- a/python/05-technical-use-cases/rag/agentic-rag/adaptive-structured-rag/src/tools/knowledge_base_tool.py +++ b/python/05-technical-use-cases/rag/agentic-rag/adaptive-structured-rag/src/tools/knowledge_base_tool.py @@ -1,8 +1,12 @@ """ Knowledge Base Tool for retrieving schema information. + +Supports both VECTOR and MANAGED knowledge base types. +Set KNOWLEDGE_BASE_TYPE environment variable to "MANAGED" to use managed search configuration. """ from strands import tool import boto3 +from botocore.config import Config import logging import os import json @@ -10,6 +14,9 @@ logger = logging.getLogger(__name__) +# Knowledge Base type: "VECTOR" (default) or "MANAGED" +KNOWLEDGE_BASE_TYPE = os.environ.get('KNOWLEDGE_BASE_TYPE', 'MANAGED').upper() + # Store the schema information for fallback WEALTH_MANAGEMENT_SCHEMA = [ { @@ -183,24 +190,74 @@ def get_schema(flag: bool = False, table_name: str = None) -> str: return _format_schema_from_data(WEALTH_MANAGEMENT_SCHEMA, table_name) # Create Bedrock client - logger.debug(f"Connecting to knowledge base: {knowledge_base_id}") - bedrock_client = boto3.client('bedrock-agent-runtime', region_name=config['aws_region']) - + logger.debug(f"Connecting to knowledge base: {knowledge_base_id} (type: {KNOWLEDGE_BASE_TYPE})") + bedrock_client = boto3.client( + 'bedrock-agent-runtime', + region_name=config['aws_region'], + config=Config(user_agent_extra='strands-samples/bedrock-kb'), + ) + # Prepare the query query = f"Describe the schema for {table_name} table" if table_name else "Describe all tables and their schemas" logger.debug(f"Querying knowledge base with: {query}") - + + # Build retrieval configuration based on KB type + if KNOWLEDGE_BASE_TYPE == "MANAGED": + # Toggle: USE_AGENTIC_RETRIEVAL=false to disable, GENERATE_RESPONSE=true for answer generation + generate_response = os.environ.get('GENERATE_RESPONSE', 'false').lower() == 'true' + use_agentic = os.environ.get('USE_AGENTIC_RETRIEVAL', 'true').lower() == 'true' + # Primary: AgenticRetrieveStream (query decomposition + managed reranking) + if use_agentic: + try: + response = bedrock_client.agentic_retrieve_stream( + messages=[{"content": {"text": query}, "role": "user"}], + retrievers=[{ + "configuration": { + "knowledgeBase": { + "knowledgeBaseId": knowledge_base_id, + "retrievalOverrides": {"maxNumberOfResults": 5}, + } + } + }], + agenticRetrieveConfiguration={ + "foundationModelType": "MANAGED", + "rerankingModelType": "MANAGED", + }, + generateResponse=generate_response, + ) + schema_info = "" + for event in response.get("stream", []): + if "result" in event: + for result in event["result"].get("results", []): + if 'content' in result and 'text' in result['content']: + schema_info += result['content']['text'] + "\n\n" + if schema_info: + logger.info("Successfully retrieved schema via AgenticRetrieveStream") + return schema_info + except Exception as e: + logger.warning(f"AgenticRetrieveStream unavailable ({e}), falling back to Retrieve") + + # Fallback: Retrieve with managedSearchConfiguration + retrieval_configuration = { + 'managedSearchConfiguration': { + 'numberOfResults': 5 + } + } + else: + # Vector KBs use vectorSearchConfiguration (default) + retrieval_configuration = { + 'vectorSearchConfiguration': { + 'numberOfResults': 5 + } + } + # Query the knowledge base response = bedrock_client.retrieve( knowledgeBaseId=knowledge_base_id, retrievalQuery={ 'text': query }, - retrievalConfiguration={ - 'vectorSearchConfiguration': { - 'numberOfResults': 5 - } - } + retrievalConfiguration=retrieval_configuration ) # Process and format the response