seahaven-slack-bot/lib/constructs/knowledge-base.ts

import * as cdk from 'aws-cdk-lib';
import { Construct } from 'constructs';
import * as s3 from 'aws-cdk-lib/aws-s3';
import * as oss from 'aws-cdk-lib/aws-opensearchserverless';
import { bedrock } from '@cdklabs/generative-ai-cdk-constructs';

export interface KnowledgeBaseProps {
  accountId: string;
  region: string;
}

export class KnowledgeBaseConstruct extends Construct {
  public readonly knowledgeBase: bedrock.VectorKnowledgeBase;
  public readonly dataSource: bedrock.S3DataSource;
  public readonly docsBucket: s3.Bucket;

  constructor(scope: Construct, id: string, props: KnowledgeBaseProps) {
    super(scope, id);

    // S3 bucket for SA8000 docs, SOPs, employee handbook, work orders, and purchase orders
    this.docsBucket = new s3.Bucket(this, 'DocsBucket', {
      bucketName: `seahaven-kb-docs-${props.accountId}`,
      encryption: s3.BucketEncryption.S3_MANAGED,
      blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
      versioned: true,
      removalPolicy: cdk.RemovalPolicy.RETAIN,
    });

    // VectorKnowledgeBase from @cdklabs/generative-ai-cdk-constructs handles:
    // - AOSS collection with correct encryption/network/access policies
    // - IAM execution role
    // - kNN vector index creation via built-in custom resource
    // - Bedrock Knowledge Base creation once the index is ready
    this.knowledgeBase = new bedrock.VectorKnowledgeBase(this, 'KnowledgeBase', {
      embeddingsModel: bedrock.BedrockFoundationModel.TITAN_EMBED_TEXT_V2_1024,
      name: 'seahaven-kb',
      description: 'SA8000 compliance docs, SOPs, employee handbook, work orders, and purchase orders',
      instruction: 'Use this knowledge base to answer questions about Sea Haven Industries company policies, SOPs, SA8000 social accountability compliance requirements, approved vendor lists, the employee handbook, work order status and history, and purchase order details.',
    });

    // Lock the auto-created AOSS network policy to private (INFRA-92).
    // @cdklabs/generative-ai-cdk-constructs hardcodes AllowFromPublic: true on the
    // VectorCollection's network policy and exposes no prop to change it, so we reach
    // the underlying CfnSecurityPolicy via the construct tree and override its Policy.
    // SourceServices: ['bedrock.amazonaws.com'] is REQUIRED — it is what keeps
    // Bedrock-managed retrieval working once public access is removed. A SourceVPCEs-only
    // policy returns 401 for Bedrock retrieval. Dashboard rule kept for console access
    // (AWS services cannot reach Dashboards regardless).
    const vectorCollection = this.knowledgeBase.vectorStore as Construct;
    const networkPolicy = vectorCollection.node.findChild('NetworkPolicy');
    if (!(networkPolicy instanceof oss.CfnSecurityPolicy)) {
      throw new Error(
        "Expected child 'NetworkPolicy' of the AOSS VectorCollection to be a CfnSecurityPolicy. " +
          'The @cdklabs/generative-ai-cdk-constructs internals may have changed — review knowledge-base.ts (INFRA-92).',
      );
    }
    const collectionName = (this.knowledgeBase.vectorStore as { collectionName?: string }).collectionName;
    if (!collectionName) {
      throw new Error(
        'Could not resolve the AOSS collection name from vectorStore — check the @cdklabs construct API (INFRA-92).',
      );
    }
    // Policy is typed as a JSON string on the L1 CfnSecurityPolicy, so it must be stringified.
    networkPolicy.addPropertyOverride(
      'Policy',
      JSON.stringify([
        {
          Rules: [
            {
              ResourceType: 'collection',
              Resource: [`collection/${collectionName}`],
            },
          ],
          AllowFromPublic: false,
          SourceServices: ['bedrock.amazonaws.com'],
        },
        {
          Rules: [
            {
              ResourceType: 'dashboard',
              Resource: [`collection/${collectionName}`],
            },
          ],
          // INFRA-92: dashboard endpoint intentionally left public for console access.
          // AWS services (incl. Bedrock) cannot reach Dashboards regardless, so this does
          // not affect the data plane. Remove this rule to fully lock down console access.
          AllowFromPublic: true,
        },
      ]),
    );

    // S3 data source — chunking configured via ChunkingStrategy.fixedSize()
    this.dataSource = new bedrock.S3DataSource(this, 'S3DataSource', {
      bucket: this.docsBucket,
      knowledgeBase: this.knowledgeBase,
      dataSourceName: 'seahaven-s3-docs',
      chunkingStrategy: bedrock.ChunkingStrategy.fixedSize({
        maxTokens: 512,
        overlapPercentage: 20,
      }),
    });
  }
}