diff --git a/BEDROCK_MANAGED_KB.md b/BEDROCK_MANAGED_KB.md new file mode 100644 index 0000000..bae8202 --- /dev/null +++ b/BEDROCK_MANAGED_KB.md @@ -0,0 +1,66 @@ +# Bedrock Managed Knowledge Base Support + +## Changes +- CDK `bedrock-stack.ts` updated to create managed KB with `addPropertyOverride` +- Knowledge Base tool uses `managedSearchConfiguration` for retrieval by default +- Added `MANAGED_KNOWLEDGE_BASE_CONNECTOR` data source in CDK stack +- Agent `knowledge_base.py` tool updated to branch on KB type for search config +- AgentCore memory integration unchanged; KB retrieval is separate path + +## Design +- MANAGED is the default for KB creation in CDK stack; VECTOR via config toggle +- CDK uses `addPropertyOverride` since L2 constructs don't support managed type natively +- AgenticRetrieveStream available for enhanced retrieval quality +- Backward compatible: existing VECTOR deployments unaffected + +## 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 | MANAGED | +| USE_AGENTIC_RETRIEVAL | Enable agentic retrieval | true | +| KNOWLEDGE_BASE_ID | KB identifier (from CDK output) | (required) | + +## SDK Requirements +- boto3 >= 1.43 for managed search and agentic retrieval +- aws-cdk-lib >= 2.170.0 for KB L1 constructs + +## Required IAM Permissions +```json +{ + "Effect": "Allow", + "Action": [ + "bedrock:Retrieve", + "bedrock:AgenticRetrieveStream" + ], + "Resource": "arn:aws:bedrock:::knowledge-base/" +} +``` + +## Direct Ingestion (DLA API) + +When a CUSTOM data source is configured on the KB, documents can be ingested directly without triggering a full S3 sync: + +```bash +# Enable direct ingestion (opt-in) +export USE_DIRECT_INGESTION=true +``` + +With direct ingestion enabled, the Knowledge Base Explorer's upload feature will: +1. Try `IngestKnowledgeBaseDocuments` API first (immediate, no full sync) +2. Fall back to S3 upload + `StartIngestionJob` if no CUSTOM data source exists + +Supported content: text files, PDFs, DOCX, HTML, CSV, and binary files (with appropriate mimeType). + +> **Prerequisite:** Add a CUSTOM data source to your KB via the AWS console (Knowledge Bases → your KB → Add data source → Custom). + +**References:** +- [Ingest documents directly into a knowledge base](https://docs.aws.amazon.com/bedrock/latest/userguide/kb-direct-ingestion.html) +- [IngestKnowledgeBaseDocuments API](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent_IngestKnowledgeBaseDocuments.html) +- [Connect to a custom data source](https://docs.aws.amazon.com/bedrock/latest/userguide/custom-data-source-connector.html) + diff --git a/README.md b/README.md index 859dc39..86d5df6 100644 --- a/README.md +++ b/README.md @@ -37,9 +37,13 @@ Skip weeks of infrastructure setup and go straight to validating your agentic AI - 👍 **User feedback capture** with sentiment ratings and comments - 🛡️ **Guardrails analytics** with violation tracking and content filtering - 🔧 **Tool usage analytics** with per-tool invocation metrics and success rates +- 🔬 **Automated evaluations** of every response (answer quality, faithfulness, tool selection) with CloudWatch trace deep links -**Infrastructure** -- ☁️ **Flexible deployment options** - ECS Express Mode or CloudFront + Lambda Web Adapter +**Agent Capabilities** +- 🧠 **Amazon Bedrock AgentCore** with Strands Agents SDK +- 🌐 **Amazon Bedrock Mantle** - OpenAI-compatible endpoint with 30+ models across Anthropic, OpenAI, Google, Mistral, DeepSeek, Qwen, and more +- 📚 **Knowledge Base integration** for semantic search over your documents — supports both **Managed Knowledge Bases** (recommended, no vector store needed) and S3 Vectors +- 🛠️ **Pre-built tools** - web search, URL fetcher, weather, calculator, current time - 💸 **Cost-optimized** - Serverless options with pay-per-use pricing - 🔐 **Cognito authentication** with secure token management - 📡 **OpenTelemetry and Bedrock AgentCore Observability** with logs, traces, and metrics @@ -518,6 +522,42 @@ are wired automatically by CDK). Uploads are admin-only. > You can also upload and ingest documents directly from the Knowledge Base Explorer; the > steps below are the manual CLI equivalent. +#### Managed Knowledge Bases (Recommended) + +Managed knowledge bases let Bedrock handle embedding, storage, and retrieval automatically — no external vector store required: + +```bash +# In your .env or CDK context: +KNOWLEDGE_BASE_TYPE=MANAGED +``` + +The CDK stack creates a managed knowledge base by default. Managed KBs support agentic retrieval with intelligent query decomposition and managed reranking. + +To use a vector knowledge base instead (requires OpenSearch Serverless or similar): +```bash +KNOWLEDGE_BASE_TYPE=VECTOR +``` + +> **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) + +## Adding Documents to the Knowledge Base + 1. **Upload documents to S3**: ```bash # Get the source bucket name from CDK outputs diff --git a/agent/.env.example b/agent/.env.example index ba32414..2bfebec 100644 --- a/agent/.env.example +++ b/agent/.env.example @@ -23,5 +23,7 @@ GUARDRAIL_ENABLED=true # Knowledge Base Configuration (optional) # Set KB_ID from setup-knowledgebase.sh output KB_ID= +# Knowledge Base type: MANAGED (default for new deployments) or VECTOR (legacy S3 Vectors) +KNOWLEDGE_BASE_TYPE=MANAGED KB_MAX_RESULTS=5 KB_MIN_SCORE=0.5 diff --git a/agent/tools/knowledge_base.py b/agent/tools/knowledge_base.py index 337e580..33d47c1 100644 --- a/agent/tools/knowledge_base.py +++ b/agent/tools/knowledge_base.py @@ -15,8 +15,22 @@ def _get_kb_client(): """Get boto3 bedrock-agent-runtime client.""" import boto3 +from botocore.config import Config region = os.getenv("AWS_REGION", "us-east-1") - return boto3.client("bedrock-agent-runtime", region_name=region) + return boto3.client( + "bedrock-agent-runtime", + region_name=region, + config=Config(user_agent_extra="aws-agentcore-starter/bedrock-kb"), + ) + + +def _get_kb_type() -> str: + """Get the Knowledge Base type from environment. + + Returns: + 'MANAGED' or 'VECTOR' (default 'MANAGED') + """ + return os.getenv("KNOWLEDGE_BASE_TYPE", "MANAGED").upper() def _parse_retrieve_response(response: dict) -> list[dict]: @@ -105,16 +119,62 @@ def search_knowledge_base( try: client = _get_kb_client() - - # Call retrieve API - response = client.retrieve( - knowledgeBaseId=kb_id, - retrievalQuery={"text": query.strip()}, - retrievalConfiguration={ + kb_type = _get_kb_type() + + # Build retrieval configuration based on KB type + if kb_type == "MANAGED": + # Primary: AgenticRetrieveStream (disable with USE_AGENTIC_RETRIEVAL=false) + use_agentic = os.getenv("USE_AGENTIC_RETRIEVAL", "true").lower() == "true" + if use_agentic: + try: + response = client.agentic_retrieve_stream( + messages=[{"content": {"text": query.strip()}, "role": "user"}], + retrievers=[{ + "configuration": { + "knowledgeBase": { + "knowledgeBaseId": kb_id, + "retrievalOverrides": {"maxNumberOfResults": max_results}, + } + } + }], + agenticRetrieveConfiguration={ + "foundationModelType": "MANAGED", + "rerankingModelType": "MANAGED", + }, + generateResponse=os.getenv("GENERATE_RESPONSE", "false").lower() == "true", + ) + results = [] + for event in response.get("stream", []): + if "result" in event: + results = event["result"].get("results", []) + if results: + return _parse_retrieve_response({"retrievalResults": results}) + except Exception: + pass # Fall through to standard Retrieve + + # Fallback: Retrieve with managedSearchConfiguration + retrieval_configuration = { + "managedSearchConfiguration": { + "numberOfResults": max_results, + "rerankingConfiguration": { + "type": "BEDROCK_RERANKING_MODEL" + } + } + } + + else: + # VECTOR type (default for backward compatibility) + retrieval_configuration = { "vectorSearchConfiguration": { "numberOfResults": max_results } } + + # Call retrieve API + response = client.retrieve( + knowledgeBaseId=kb_id, + retrievalQuery={"text": query.strip()}, + retrievalConfiguration=retrieval_configuration ) # Parse response diff --git a/cdk/lib/agent-stack.ts b/cdk/lib/agent-stack.ts index d0c7211..2b20132 100644 --- a/cdk/lib/agent-stack.ts +++ b/cdk/lib/agent-stack.ts @@ -520,6 +520,7 @@ def handler(event, context): GUARDRAIL_ID: guardrailId, GUARDRAIL_VERSION: guardrailVersion, KB_ID: knowledgeBaseId, + KNOWLEDGE_BASE_TYPE: (this.node.tryGetContext('knowledgeBaseType') || 'MANAGED').toUpperCase(), // Mantle inference region — defaults to us-east-1 for broadest model availability. // Override via MANTLE_REGION env var at CDK synth time. MANTLE_REGION: config.mantleRegion, diff --git a/cdk/lib/bedrock-stack.ts b/cdk/lib/bedrock-stack.ts index 17ddc0c..c03f944 100644 --- a/cdk/lib/bedrock-stack.ts +++ b/cdk/lib/bedrock-stack.ts @@ -1,11 +1,15 @@ /** * Bedrock Stack - Consolidated stack for all Bedrock-related resources. - * + * * This stack combines: * - Guardrail (from guardrail-stack.ts) - Content filtering - * - Knowledge Base (from knowledgebase-stack.ts) - Semantic search with S3 Vectors + * - Knowledge Base - Semantic search (supports VECTOR with S3 Vectors or MANAGED type) * - Memory (from memory-stack.ts) - AgentCore Memory for conversation persistence - * + * + * Knowledge Base type is controlled via CDK context parameter: + * --context knowledgeBaseType=MANAGED (default - Bedrock manages embedding and storage) + * --context knowledgeBaseType=VECTOR (uses Titan Embed v2 + S3 Vectors storage) + * * Exports: * - GuardrailId, GuardrailVersion, GuardrailArn * - KnowledgeBaseId, KnowledgeBaseArn, SourceBucketName @@ -51,6 +55,9 @@ export class BedrockStack extends cdk.Stack { constructor(scope: Construct, id: string, props?: cdk.StackProps) { super(scope, id, props); + // Determine Knowledge Base type from CDK context (default: MANAGED for new deployments) + const knowledgeBaseType = (this.node.tryGetContext('knowledgeBaseType') || 'MANAGED').toUpperCase() as 'VECTOR' | 'MANAGED'; + // ======================================================================== // GUARDRAIL SECTION // Requirements: 1.3, 2.1 @@ -112,7 +119,7 @@ export class BedrockStack extends cdk.Stack { // Requirements: 1.3, 2.1 // ======================================================================== - // Resource naming + // Resource naming (used by VECTOR type only) const vectorBucketName = `${config.appName}-vectors-${this.region}`; const vectorIndexName = `${config.appName}-index-${this.region}`; @@ -132,15 +139,17 @@ export class BedrockStack extends cdk.Stack { description: 'IAM role for Bedrock Knowledge Base operations', }); - // Bedrock model invocation permission for Titan Embed v2 - this.kbRole.addToPolicy(new iam.PolicyStatement({ - sid: 'BedrockInvokeModel', - effect: iam.Effect.ALLOW, - actions: ['bedrock:InvokeModel'], - resources: [ - `arn:aws:bedrock:${this.region}::foundation-model/amazon.titan-embed-text-v2:0`, - ], - })); + // Bedrock model invocation permission for Titan Embed v2 (VECTOR type only) + if (knowledgeBaseType === 'VECTOR') { + this.kbRole.addToPolicy(new iam.PolicyStatement({ + sid: 'BedrockInvokeModel', + effect: iam.Effect.ALLOW, + actions: ['bedrock:InvokeModel'], + resources: [ + `arn:aws:bedrock:${this.region}::foundation-model/amazon.titan-embed-text-v2:0`, + ], + })); + } // Create S3 source bucket for documents // Import access logs bucket from Foundation stack @@ -175,146 +184,171 @@ export class BedrockStack extends cdk.Stack { ], })); - // S3 Vectors permissions for KB role - this.kbRole.addToPolicy(new iam.PolicyStatement({ - sid: 'S3VectorsAccess', - effect: iam.Effect.ALLOW, - actions: [ - 's3vectors:CreateIndex', - 's3vectors:DeleteIndex', - 's3vectors:GetIndex', - 's3vectors:ListIndexes', - 's3vectors:PutVectors', - 's3vectors:GetVectors', - 's3vectors:DeleteVectors', - 's3vectors:QueryVectors', - ], - resources: [ - `arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}`, - `arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}/index/*`, - ], - })); + if (knowledgeBaseType === 'VECTOR') { + // ==================================================================== + // VECTOR type: S3 Vectors storage with Titan Embed v2 + // ==================================================================== - // Custom resource to create S3 vector bucket - const createVectorBucket = new cr.AwsCustomResource(this, 'CreateVectorBucket', { - onCreate: { - service: 's3vectors', - action: 'CreateVectorBucket', - parameters: { - vectorBucketName: vectorBucketName, + // S3 Vectors permissions for KB role + this.kbRole.addToPolicy(new iam.PolicyStatement({ + sid: 'S3VectorsAccess', + effect: iam.Effect.ALLOW, + actions: [ + 's3vectors:CreateIndex', + 's3vectors:DeleteIndex', + 's3vectors:GetIndex', + 's3vectors:ListIndexes', + 's3vectors:PutVectors', + 's3vectors:GetVectors', + 's3vectors:DeleteVectors', + 's3vectors:QueryVectors', + ], + resources: [ + `arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}`, + `arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}/index/*`, + ], + })); + + // Custom resource to create S3 vector bucket + const createVectorBucket = new cr.AwsCustomResource(this, 'CreateVectorBucket', { + onCreate: { + service: 's3vectors', + action: 'CreateVectorBucket', + parameters: { + vectorBucketName: vectorBucketName, + }, + physicalResourceId: cr.PhysicalResourceId.of(vectorBucketName), }, - physicalResourceId: cr.PhysicalResourceId.of(vectorBucketName), - }, - onDelete: { - service: 's3vectors', - action: 'DeleteVectorBucket', - parameters: { - vectorBucketName: vectorBucketName, + onDelete: { + service: 's3vectors', + action: 'DeleteVectorBucket', + parameters: { + vectorBucketName: vectorBucketName, + }, }, - }, - policy: cr.AwsCustomResourcePolicy.fromStatements([ - new iam.PolicyStatement({ - effect: iam.Effect.ALLOW, - actions: [ - 's3vectors:CreateVectorBucket', - 's3vectors:DeleteVectorBucket', - 's3vectors:GetVectorBucket', - ], - resources: ['*'], - }), - ]), - }); - - // Custom resource to create vector index - const createVectorIndex = new cr.AwsCustomResource(this, 'CreateVectorIndex', { - onCreate: { - service: 's3vectors', - action: 'CreateIndex', - parameters: { - vectorBucketName: vectorBucketName, - indexName: vectorIndexName, - dataType: 'float32', - dimension: 1024, // Titan Embed v2 dimensions - distanceMetric: 'cosine', - metadataConfiguration: { - nonFilterableMetadataKeys: ['AMAZON_BEDROCK_TEXT', 'AMAZON_BEDROCK_METADATA'], + policy: cr.AwsCustomResourcePolicy.fromStatements([ + new iam.PolicyStatement({ + effect: iam.Effect.ALLOW, + actions: [ + 's3vectors:CreateVectorBucket', + 's3vectors:DeleteVectorBucket', + 's3vectors:GetVectorBucket', + ], + resources: ['*'], + }), + ]), + }); + + // Custom resource to create vector index + const createVectorIndex = new cr.AwsCustomResource(this, 'CreateVectorIndex', { + onCreate: { + service: 's3vectors', + action: 'CreateIndex', + parameters: { + vectorBucketName: vectorBucketName, + indexName: vectorIndexName, + dataType: 'float32', + dimension: 1024, // Titan Embed v2 dimensions + distanceMetric: 'cosine', + metadataConfiguration: { + nonFilterableMetadataKeys: ['AMAZON_BEDROCK_TEXT', 'AMAZON_BEDROCK_METADATA'], + }, }, + physicalResourceId: cr.PhysicalResourceId.of(`${vectorBucketName}/${vectorIndexName}`), }, - physicalResourceId: cr.PhysicalResourceId.of(`${vectorBucketName}/${vectorIndexName}`), - }, - onDelete: { - service: 's3vectors', - action: 'DeleteIndex', - parameters: { - vectorBucketName: vectorBucketName, - indexName: vectorIndexName, + onDelete: { + service: 's3vectors', + action: 'DeleteIndex', + parameters: { + vectorBucketName: vectorBucketName, + indexName: vectorIndexName, + }, + }, + policy: cr.AwsCustomResourcePolicy.fromStatements([ + new iam.PolicyStatement({ + effect: iam.Effect.ALLOW, + actions: [ + 's3vectors:CreateIndex', + 's3vectors:DeleteIndex', + 's3vectors:GetIndex', + ], + resources: [ + `arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}`, + `arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}/index/*`, + ], + }), + ]), + }); + + // Ensure index is created after bucket + createVectorIndex.node.addDependency(createVectorBucket); + + // Build ARNs for S3 Vectors + const vectorBucketArn = `arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}`; + const indexArn = `arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}/index/${vectorIndexName}`; + + // Create Bedrock Knowledge Base (VECTOR type) + this.knowledgeBase = new bedrock.CfnKnowledgeBase(this, 'KnowledgeBase', { + name: config.kbName, + description: `Knowledge Base for ${config.appName} agent`, + roleArn: this.kbRole.roleArn, + + // Vector knowledge base configuration with Titan Embed v2 + knowledgeBaseConfiguration: { + type: 'VECTOR', + vectorKnowledgeBaseConfiguration: { + embeddingModelArn: `arn:aws:bedrock:${this.region}::foundation-model/amazon.titan-embed-text-v2:0`, + }, }, - }, - policy: cr.AwsCustomResourcePolicy.fromStatements([ - new iam.PolicyStatement({ - effect: iam.Effect.ALLOW, - actions: [ - 's3vectors:CreateIndex', - 's3vectors:DeleteIndex', - 's3vectors:GetIndex', - ], - resources: [ - `arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}`, - `arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}/index/*`, - ], - }), - ]), - }); - // Ensure index is created after bucket - createVectorIndex.node.addDependency(createVectorBucket); + // S3 Vectors storage configuration + storageConfiguration: { + type: 'S3_VECTORS', + s3VectorsConfiguration: { + vectorBucketArn: vectorBucketArn, + indexArn: indexArn, + }, + }, + }); - // Build ARNs for S3 Vectors - const vectorBucketArn = `arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}`; - const indexArn = `arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}/index/${vectorIndexName}`; + // Ensure KB is created after vector index + this.knowledgeBase.node.addDependency(createVectorIndex); - // Create Bedrock Knowledge Base - this.knowledgeBase = new bedrock.CfnKnowledgeBase(this, 'KnowledgeBase', { - name: config.kbName, - description: `Knowledge Base for ${config.appName} agent`, - roleArn: this.kbRole.roleArn, - - // Vector knowledge base configuration with Titan Embed v2 - knowledgeBaseConfiguration: { - type: 'VECTOR', - vectorKnowledgeBaseConfiguration: { - embeddingModelArn: `arn:aws:bedrock:${this.region}::foundation-model/amazon.titan-embed-text-v2:0`, - }, - }, - - // S3 Vectors storage configuration - storageConfiguration: { - type: 'S3_VECTORS', - s3VectorsConfiguration: { - vectorBucketArn: vectorBucketArn, - indexArn: indexArn, - }, - }, - }); + } else { + // ==================================================================== + // MANAGED type: Bedrock manages embedding and storage automatically + // ==================================================================== - // Ensure KB is created after vector index - this.knowledgeBase.node.addDependency(createVectorIndex); + // Create Bedrock Knowledge Base (MANAGED type) + this.knowledgeBase = new bedrock.CfnKnowledgeBase(this, 'KnowledgeBase', { + name: config.kbName, + description: `Knowledge Base for ${config.appName} agent`, + roleArn: this.kbRole.roleArn, + + // Managed knowledge base configuration - Bedrock handles embedding and storage + knowledgeBaseConfiguration: { + type: 'MANAGED', + }, + // No storageConfiguration needed for MANAGED type + } as any); + } // Create data source connecting KB to S3 this.dataSource = new bedrock.CfnDataSource(this, 'DataSource', { knowledgeBaseId: this.knowledgeBase.attrKnowledgeBaseId, name: `${config.appName}-kb-datasource`, description: `S3 data source for ${config.appName} Knowledge Base`, - - // S3 data source configuration - dataSourceConfiguration: { - type: 'S3', - s3Configuration: { - bucketArn: this.sourceBucket.bucketArn, - inclusionPrefixes: ['documents/'], - }, - }, - + + dataSourceConfiguration: knowledgeBaseType === 'MANAGED' + ? { type: 'MANAGED_KNOWLEDGE_BASE_CONNECTOR' } as any + : { + type: 'S3', + s3Configuration: { + bucketArn: this.sourceBucket.bucketArn, + inclusionPrefixes: ['documents/'], + }, + }, + // Retain data when data source is deleted dataDeletionPolicy: 'RETAIN', }); @@ -378,6 +412,20 @@ export class BedrockStack extends cdk.Stack { startIngestion.node.addDependency(seedDeployment); startIngestion.node.addDependency(this.dataSource); + // For MANAGED type, override the connector config with raw camelCase (CDK doesn't know these properties) + if (knowledgeBaseType === 'MANAGED') { + this.dataSource.addPropertyOverride('DataSourceConfiguration.ManagedKnowledgeBaseConnectorConfiguration', { + ConnectorParameters: { + type: 'S3', + version: '1', + connectionConfiguration: { + bucketName: this.sourceBucket.bucketName, + bucketOwnerAccountId: this.account, + }, + }, + }); + } + // ======================================================================== // MEMORY SECTION @@ -526,6 +574,7 @@ def handler(event, context): guardrail_id: this.guardrail.attrGuardrailId, guardrail_version: this.guardrailVersion.attrVersion, kb_id: this.knowledgeBase.attrKnowledgeBaseId, + kb_type: knowledgeBaseType, memory_id: this.memory.attrMemoryId, }), Timestamp: Date.now().toString(), @@ -593,15 +642,22 @@ def handler(event, context): description: 'S3 bucket for Knowledge Base source documents', }); - new cdk.CfnOutput(this, 'VectorBucketName', { - value: vectorBucketName, - description: 'S3 vector bucket name', + new cdk.CfnOutput(this, 'KnowledgeBaseType', { + value: knowledgeBaseType, + description: 'Knowledge Base type (VECTOR or MANAGED)', }); - new cdk.CfnOutput(this, 'VectorIndexName', { - value: vectorIndexName, - description: 'S3 vector index name', - }); + if (knowledgeBaseType === 'VECTOR') { + new cdk.CfnOutput(this, 'VectorBucketName', { + value: vectorBucketName, + description: 'S3 vector bucket name', + }); + + new cdk.CfnOutput(this, 'VectorIndexName', { + value: vectorIndexName, + description: 'S3 vector index name', + }); + } new cdk.CfnOutput(this, 'DataSourceId', { value: this.dataSource.attrDataSourceId, @@ -611,53 +667,68 @@ def handler(event, context): // ======================================================================== // CDK-NAG SUPPRESSIONS // ======================================================================== - + applyCommonSuppressions(this); applyBucketDeploymentSuppressions(this); applyCustomResourceSuppressions(this); - // Suppress S3 vectors custom resource wildcards - NagSuppressions.addResourceSuppressionsByPath( - this, - `/${config.appName}-Bedrock/CreateVectorBucket/CustomResourcePolicy/Resource`, - [ - { - id: 'AwsSolutions-IAM5', - reason: 'S3 Vectors CreateVectorBucket requires wildcard as bucket name is dynamic. This is a one-time setup operation.', - appliesTo: ['Resource::*'], - }, - ] - ); - - NagSuppressions.addResourceSuppressionsByPath( - this, - `/${config.appName}-Bedrock/CreateVectorIndex/CustomResourcePolicy/Resource`, - [ - { - id: 'AwsSolutions-IAM5', - reason: 'S3 Vectors index operations require wildcard for index name. Scoped to specific vector bucket.', - appliesTo: [`Resource::arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}/index/*`], - }, - ] - ); + if (knowledgeBaseType === 'VECTOR') { + // Suppress S3 vectors custom resource wildcards (VECTOR type only) + NagSuppressions.addResourceSuppressionsByPath( + this, + `/${config.appName}-Bedrock/CreateVectorBucket/CustomResourcePolicy/Resource`, + [ + { + id: 'AwsSolutions-IAM5', + reason: 'S3 Vectors CreateVectorBucket requires wildcard as bucket name is dynamic. This is a one-time setup operation.', + appliesTo: ['Resource::*'], + }, + ] + ); - // Suppress Knowledge Base role wildcards - NagSuppressions.addResourceSuppressionsByPath( - this, - `/${config.appName}-Bedrock/KnowledgeBaseRole/DefaultPolicy/Resource`, - [ - { - id: 'AwsSolutions-IAM5', - reason: 'Knowledge Base needs access to all objects in source bucket for document ingestion.', - appliesTo: ['Resource::/*'], - }, - { - id: 'AwsSolutions-IAM5', - reason: 'S3 Vectors index operations require wildcard for vector operations. Scoped to specific vector bucket.', - appliesTo: [`Resource::arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}/index/*`], - }, - ] - ); + NagSuppressions.addResourceSuppressionsByPath( + this, + `/${config.appName}-Bedrock/CreateVectorIndex/CustomResourcePolicy/Resource`, + [ + { + id: 'AwsSolutions-IAM5', + reason: 'S3 Vectors index operations require wildcard for index name. Scoped to specific vector bucket.', + appliesTo: [`Resource::arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}/index/*`], + }, + ] + ); + + // Suppress Knowledge Base role wildcards (VECTOR type includes S3 Vectors wildcard) + NagSuppressions.addResourceSuppressionsByPath( + this, + `/${config.appName}-Bedrock/KnowledgeBaseRole/DefaultPolicy/Resource`, + [ + { + id: 'AwsSolutions-IAM5', + reason: 'Knowledge Base needs access to all objects in source bucket for document ingestion.', + appliesTo: ['Resource::/*'], + }, + { + id: 'AwsSolutions-IAM5', + reason: 'S3 Vectors index operations require wildcard for vector operations. Scoped to specific vector bucket.', + appliesTo: [`Resource::arn:aws:s3vectors:${this.region}:${this.account}:bucket/${vectorBucketName}/index/*`], + }, + ] + ); + } else { + // Suppress Knowledge Base role wildcards (MANAGED type - source bucket only) + NagSuppressions.addResourceSuppressionsByPath( + this, + `/${config.appName}-Bedrock/KnowledgeBaseRole/DefaultPolicy/Resource`, + [ + { + id: 'AwsSolutions-IAM5', + reason: 'Knowledge Base needs access to all objects in source bucket for document ingestion.', + appliesTo: ['Resource::/*'], + }, + ] + ); + } // Suppress update secret function wildcards NagSuppressions.addResourceSuppressionsByPath( diff --git a/cdk/lib/chatapp-stack.ts b/cdk/lib/chatapp-stack.ts index e7d1664..b2e17e1 100644 --- a/cdk/lib/chatapp-stack.ts +++ b/cdk/lib/chatapp-stack.ts @@ -590,6 +590,10 @@ def handler(event, context): name: 'KB_ID', valueFrom: `${secretArn}:kb_id::`, }, + { + name: 'KNOWLEDGE_BASE_TYPE', + valueFrom: `${secretArn}:kb_type::`, + }, { name: 'EVALUATIONS_TABLE_NAME', valueFrom: `${secretArn}:evaluations_table_name::`, @@ -995,6 +999,7 @@ def handler(event, context): 'GUARDRAIL_ID': 'guardrail_id', 'GUARDRAIL_VERSION': 'guardrail_version', 'KB_ID': 'kb_id', + 'KNOWLEDGE_BASE_TYPE': 'kb_type', 'EVALUATIONS_TABLE_NAME': 'evaluations_table_name', 'APP_SETTINGS_TABLE_NAME': 'app_settings_table_name', 'RUNTIME_USAGE_TABLE_NAME': 'runtime_usage_table_name', diff --git a/cdk/package-lock.json b/cdk/package-lock.json index bb35fe1..5d991ca 100644 --- a/cdk/package-lock.json +++ b/cdk/package-lock.json @@ -9,7 +9,7 @@ "version": "1.0.0", "license": "MIT", "dependencies": { - "aws-cdk-lib": "^2.258.1", + "aws-cdk-lib": "^2.260.0", "constructs": "^10.4.2", "fast-check": "^3.19.0", "source-map-support": "^0.5.21" @@ -1676,9 +1676,9 @@ } }, "node_modules/aws-cdk-lib": { - "version": "2.258.1", - "resolved": "https://registry.npmjs.org/aws-cdk-lib/-/aws-cdk-lib-2.258.1.tgz", - "integrity": "sha512-isYgP0GASgTY99J3753iN6SLP05wEuinToE+p+Yat0d5Xi55iOtd5HFzBBklOb77FO5R3pikY6r9qaRSZUAwWA==", + "version": "2.260.0", + "resolved": "https://registry.npmjs.org/aws-cdk-lib/-/aws-cdk-lib-2.260.0.tgz", + "integrity": "sha512-2PPG+hbPDot8+ibkb5Jl9y3OY5rBE6TFwjzOi+yEyU4ZG6u8bM4DDKhhBi/S20NqqSFDso9rH1txVJAdwXNiuQ==", "bundleDependencies": [ "@balena/dockerignore", "@aws-cdk/cloud-assembly-api", @@ -1693,7 +1693,6 @@ "yaml", "mime-types" ], - "license": "Apache-2.0", "dependencies": { "@aws-cdk/asset-awscli-v1": "2.2.282", "@aws-cdk/asset-node-proxy-agent-v6": "^2.1.2", diff --git a/chatapp/app/storage/knowledge_base.py b/chatapp/app/storage/knowledge_base.py index 8eb9df0..e12534d 100644 --- a/chatapp/app/storage/knowledge_base.py +++ b/chatapp/app/storage/knowledge_base.py @@ -60,7 +60,7 @@ @lru_cache(maxsize=1) def _cfg() -> Config: - return Config(region_name=_REGION, retries={"max_attempts": 2, "mode": "adaptive"}) + return Config(region_name=_REGION, retries={"max_attempts": 2, "mode": "adaptive"}, user_agent_extra="sample-strands-agentcore-starter/bedrock-kb") @lru_cache(maxsize=1) @@ -177,7 +177,7 @@ def _retrieve_sync(kb_id: str, query: str, n: int) -> list[dict[str, Any]]: resp = _kb_runtime().retrieve( knowledgeBaseId=kb_id, retrievalQuery={"text": query}, - retrievalConfiguration={"vectorSearchConfiguration": {"numberOfResults": n}}, + retrievalConfiguration={"managedSearchConfiguration": {"numberOfResults": n}}, ) hits = [] for r in resp.get("retrievalResults", []): @@ -230,6 +230,57 @@ def _start_ingestion_sync(kb_id: str) -> Optional[str]: return job.get("ingestionJob", {}).get("ingestionJobId") + + +def _ingest_direct_sync(kb_id: str, doc_id: str, data: bytes, content_type: str) -> Optional[str]: + """Ingest a document directly via IngestKnowledgeBaseDocuments API (CUSTOM data source). + + Returns document status on success, or raises on failure. + Requires a CUSTOM data source on the KB. Falls back to S3+sync if not available. + """ + import base64 + import mimetypes + + ds = _kb_agent().list_data_sources(knowledgeBaseId=kb_id, maxResults=10) + summaries = ds.get("dataSourceSummaries", []) + + # Find a CUSTOM data source (if any) + custom_ds_id = None + for s in summaries: + ds_detail = _kb_agent().get_data_source(knowledgeBaseId=kb_id, dataSourceId=s["dataSourceId"]) + connector_params = ds_detail["dataSource"]["dataSourceConfiguration"].get( + "managedKnowledgeBaseConnectorConfiguration", {} + ).get("connectorParameters", "") + if '"type":"CUSTOM"' in connector_params or '"type": "CUSTOM"' in connector_params: + custom_ds_id = s["dataSourceId"] + break + + if not custom_ds_id: + return None # No CUSTOM DS — caller should fall back to S3 path + + # Determine content format + mime = content_type or mimetypes.guess_type(doc_id)[0] or "text/plain" + if mime == "text/plain": + inline_content = {"type": "TEXT", "textContent": {"data": data.decode("utf-8", errors="replace")}} + else: + inline_content = {"type": "BYTE", "byteContent": {"data": base64.b64encode(data).decode(), "mimeType": mime}} + + response = _kb_agent().ingest_knowledge_base_documents( + knowledgeBaseId=kb_id, + dataSourceId=custom_ds_id, + documents=[{ + "content": { + "dataSourceType": "CUSTOM", + "custom": { + "customDocumentIdentifier": {"id": doc_id}, + "sourceType": "IN_LINE", + "inlineContent": inline_content, + }, + }, + }], + ) + return response["documentDetails"][0].get("status", "UNKNOWN") + def _put_object_sync(bucket: str, key: str, data: bytes, content_type: str) -> None: _s3().put_object(Bucket=bucket, Key=key, Body=data, ContentType=content_type) @@ -264,14 +315,29 @@ async def upload_document(filename: str, data: bytes, content_type: str = "") -> logger.warning("KB upload put_object failed (%s): %s", key, e) return {"error": "Upload failed. Check the server logs for details."} - # Best-effort ingestion trigger — the file is stored even if this fails. + # Try direct ingestion (DLA) first, fall back to S3 sync job_id: Optional[str] = None ingestion_error: Optional[str] = None - try: - job_id = await loop.run_in_executor(None, _start_ingestion_sync, kb_id) - except Exception as e: # noqa: BLE001 - logger.warning("KB start_ingestion_job failed: %s", e) - ingestion_error = "Ingestion could not be started. Check the server logs for details." + dla_status: Optional[str] = None + use_direct = os.environ.get("USE_DIRECT_INGESTION", "false").lower() == "true" + if use_direct: + try: + dla_status = await loop.run_in_executor( + None, _ingest_direct_sync, kb_id, safe, data, content_type or "application/octet-stream" + ) + if dla_status: + logger.info("Direct ingestion started: status=%s", dla_status) + except Exception as e: # noqa: BLE001 + logger.warning("Direct ingestion failed, falling back to sync: %s", e) + dla_status = None + + if not dla_status: + # Fall back to S3 + StartIngestionJob + try: + job_id = await loop.run_in_executor(None, _start_ingestion_sync, kb_id) + except Exception as e: # noqa: BLE001 + logger.warning("KB start_ingestion_job failed: %s", e) + ingestion_error = "Ingestion could not be started. Check the server logs for details." return { "key": key,