seahaven-slack-bot/lib/constructs/bedrock-agent.ts

import * as cdk from 'aws-cdk-lib';
import { Construct } from 'constructs';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as lambdaNodejs from 'aws-cdk-lib/aws-lambda-nodejs';
import * as dynamodb from 'aws-cdk-lib/aws-dynamodb';
import * as secretsmanager from 'aws-cdk-lib/aws-secretsmanager';
import * as bedrock from 'aws-cdk-lib/aws-bedrock';
import * as path from 'path';

export interface BedrockAgentProps {
  accountId: string;
  region: string;
  knowledgeBaseId: string;
  knowledgeBaseArn: string;
}

export class BedrockAgentConstruct extends Construct {
  public readonly agent: bedrock.CfnAgent;
  public readonly agentAlias: bedrock.CfnAgentAlias;
  public readonly qboLambda: lambdaNodejs.NodejsFunction;
  public readonly mapsLambda: lambdaNodejs.NodejsFunction;
  public readonly woPoLambda: lambdaNodejs.NodejsFunction;

  // Cross-region inference profile — required for Claude 4.x on Bedrock Agents
  private static readonly MODEL_ID = 'us.anthropic.claude-sonnet-4-5-20250929-v1:0';

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

    // Reference secrets created manually in Secrets Manager (see README for structure)
    const qboSecret = secretsmanager.Secret.fromSecretNameV2(
      this, 'QBOSecret', 'seahaven/qbo/oauth',
    );
    const mapsSecret = secretsmanager.Secret.fromSecretNameV2(
      this, 'MapsSecret', 'seahaven/google/maps-api-key',
    );

    const bundling: lambdaNodejs.BundlingOptions = {
      externalModules: ['@aws-sdk/*'],
      minify: true,
      sourceMap: false,
    };

    // ── QBO vendor lookup action group Lambda ─────────────────────────────────
    this.qboLambda = new lambdaNodejs.NodejsFunction(this, 'QBOLookupFn', {
      functionName: 'seahaven-qbo-lookup',
      entry: path.join(__dirname, '../../lambda/qbo-lookup/index.ts'),
      handler: 'handler',
      runtime: lambda.Runtime.NODEJS_22_X,
      timeout: cdk.Duration.seconds(30),
      memorySize: 256,
      environment: { QBO_SECRET_ARN: qboSecret.secretArn },
      bundling,
    });
    qboSecret.grantRead(this.qboLambda);

    // ── Google Maps lookup action group Lambda ────────────────────────────────
    this.mapsLambda = new lambdaNodejs.NodejsFunction(this, 'MapsLookupFn', {
      functionName: 'seahaven-maps-lookup',
      entry: path.join(__dirname, '../../lambda/maps-lookup/index.ts'),
      handler: 'handler',
      runtime: lambda.Runtime.NODEJS_22_X,
      timeout: cdk.Duration.seconds(30),
      memorySize: 256,
      environment: { MAPS_SECRET_ARN: mapsSecret.secretArn },
      bundling,
    });
    mapsSecret.grantRead(this.mapsLambda);

    // ── WO/PO direct lookup action group Lambda ──────────────────────────────
    const workOrdersTable = dynamodb.Table.fromTableName(
      this, 'WorkOrdersTable', 'WorkOrders',
    );
    const commentsTable = dynamodb.Table.fromTableName(
      this, 'WorkOrderCommentsTable', 'WorkOrderComments',
    );
    const poTable = dynamodb.Table.fromTableName(
      this, 'PurchaseOrdersTable', 'purchase-orders',
    );

    // Site assignments table — stores Amazon facility site codes and addresses
    const sitesTable = new dynamodb.Table(this, 'SiteAssignmentsTable', {
      tableName: 'SiteAssignments',
      partitionKey: { name: 'siteCode', type: dynamodb.AttributeType.STRING },
      billingMode: dynamodb.BillingMode.PAY_PER_REQUEST,
      removalPolicy: cdk.RemovalPolicy.RETAIN,
    });
    sitesTable.addGlobalSecondaryIndex({
      indexName: 'by-state',
      partitionKey: { name: 'state', type: dynamodb.AttributeType.STRING },
      projectionType: dynamodb.ProjectionType.ALL,
    });

    this.woPoLambda = new lambdaNodejs.NodejsFunction(this, 'WoPoLookupFn', {
      functionName: 'seahaven-wo-po-lookup',
      entry: path.join(__dirname, '../../lambda/wo-po-lookup/index.ts'),
      handler: 'handler',
      runtime: lambda.Runtime.NODEJS_22_X,
      timeout: cdk.Duration.seconds(30),
      memorySize: 256,
      environment: {
        WORK_ORDERS_TABLE: workOrdersTable.tableName,
        COMMENTS_TABLE: commentsTable.tableName,
        PO_TABLE: poTable.tableName,
        SITES_TABLE: sitesTable.tableName,
      },
      bundling,
    });
    workOrdersTable.grantReadData(this.woPoLambda);
    commentsTable.grantReadData(this.woPoLambda);
    poTable.grantReadData(this.woPoLambda);
    sitesTable.grantReadData(this.woPoLambda);

    // ── Bedrock Agent execution role ──────────────────────────────────────────
    const agentRole = new iam.Role(this, 'AgentRole', {
      roleName: 'AmazonBedrockExecutionRoleForAgents_seahaven',
      assumedBy: new iam.ServicePrincipal('bedrock.amazonaws.com', {
        conditions: {
          StringEquals: { 'aws:SourceAccount': props.accountId },
          ArnLike: {
            'aws:SourceArn': `arn:aws:bedrock:${props.region}:${props.accountId}:agent/*`,
          },
        },
      }),
    });

    agentRole.addToPolicy(new iam.PolicyStatement({
      actions: ['bedrock:InvokeModel', 'bedrock:InvokeModelWithResponseStream'],
      resources: [
        // Cross-region inference profile (us.*) routes to multiple regions — wildcard region required
        `arn:aws:bedrock:*::foundation-model/anthropic.claude-sonnet-4-5-20250929-v1:0`,
        `arn:aws:bedrock:${props.region}:${props.accountId}:inference-profile/${BedrockAgentConstruct.MODEL_ID}`,
      ],
    }));

    agentRole.addToPolicy(new iam.PolicyStatement({
      actions: ['bedrock:Retrieve'],
      resources: [props.knowledgeBaseArn],
    }));

    // Allow Bedrock to invoke the action group Lambdas
    this.qboLambda.addPermission('BedrockInvokeQBO', {
      principal: new iam.ServicePrincipal('bedrock.amazonaws.com'),
      sourceAccount: props.accountId,
    });
    this.mapsLambda.addPermission('BedrockInvokeMaps', {
      principal: new iam.ServicePrincipal('bedrock.amazonaws.com'),
      sourceAccount: props.accountId,
    });
    this.woPoLambda.addPermission('BedrockInvokeWoPo', {
      principal: new iam.ServicePrincipal('bedrock.amazonaws.com'),
      sourceAccount: props.accountId,
    });

    // ── Agent instruction (system prompt) ────────────────────────────────────
    const instruction = `You are the Sea Haven Industries internal assistant, accessible to employees via Slack direct message. Sea Haven is a facility services company that provides maintenance, repair, and facility services to clients like Amazon.

IMPORTANT: You work FOR Sea Haven. Sea Haven is our company — we are the vendor/contractor that gets dispatched to job sites. When work orders or comments mention "Sea Haven" being dispatched or assigned, that means OUR team was sent. Never suggest Sea Haven as an external vendor to contact — we ARE Sea Haven. When users ask for a "local vendor," they mean a subcontractor or specialty trade vendor to handle work on our behalf.

You help employees with:
1. Finding vendors and contractors for facility work
2. Company policies, SOPs, and SA8000 social accountability compliance questions
3. Employee handbook questions
4. Work order lookups — status, history, assignments, comments, and details for any work order by its WO number
5. Purchase order lookups — status, line items, suppliers, ship-to details, and dates for any PO by its PO number
6. Amazon site lookups — address, location, and details for any Amazon facility by its site code (e.g., "ABE2", "DFW6"), or listing all sites in a given state

## Vendor Query Rules — STRICTLY follow this priority order:

Step 1: ALWAYS call QBO_Lookup.search_vendors first. This searches our QuickBooks Online account for existing, vetted vendors we already have a relationship with.

Step 2: If QBO_Lookup returns no suitable match, search the knowledge base for approved vendor documentation or lists.

Step 3: ONLY if both QBO and the knowledge base return nothing suitable should you call Google_Maps_Lookup.search_nearby_vendors to find new options.

When presenting vendor results:
- QBO vendors: include name, trade/specialty, phone, email, and last updated date
- Knowledge base vendors: cite the source document
- Google Maps vendors: clearly label these as NEW (not yet vetted), include name, address, phone, and rating

## Work Order & Purchase Order Queries:
When a user asks about a work order (WO) or purchase order (PO) by number, ALWAYS use the WO_PO_Lookup action group to retrieve the record directly. Do NOT use the knowledge base for WO/PO lookups by number — the action group queries the database directly and is more reliable. Present the returned data clearly and concisely. The output is already well-structured — relay the key information without adding excessive formatting or repeating section headers verbatim. Summarize the current status and most recent updates first, then include the full comment history.

## Site Lookups:
When a user asks about an Amazon site by its code (e.g., "ABE2", "DFW6", "WWY1"), ALWAYS use the WO_PO_Lookup.lookup_site function. You can also list all sites in a state by passing the state abbreviation. Do NOT use the knowledge base for site code lookups — the action group queries the database directly and is more reliable.

## General Questions:
Use the knowledge base for policy, SOP, SA8000 compliance, and handbook questions. Cite the specific document or section when possible. If the information is not in the knowledge base, say so clearly — do not guess.

Keep responses concise, professional, and actionable.`;

    // ── CfnAgent ──────────────────────────────────────────────────────────────
    this.agent = new bedrock.CfnAgent(this, 'Agent', {
      agentName: 'seahaven-assistant',
      description: 'Sea Haven Industries internal Slack assistant',
      agentResourceRoleArn: agentRole.roleArn,
      foundationModel: BedrockAgentConstruct.MODEL_ID,
      instruction,
      idleSessionTtlInSeconds: 1800, // 30 min — matches the processor's session window
      knowledgeBases: [
        {
          knowledgeBaseId: props.knowledgeBaseId,
          description: 'Sea Haven internal documents: SOPs, SA8000 compliance docs, employee handbook, approved vendor lists, work orders, and purchase orders',
          knowledgeBaseState: 'ENABLED',
        },
      ],
      actionGroups: [
        {
          actionGroupName: 'QBO_Lookup',
          description: 'Search QuickBooks Online for existing Sea Haven vendors by trade or name',
          actionGroupState: 'ENABLED',
          actionGroupExecutor: { lambda: this.qboLambda.functionArn },
          functionSchema: {
            functions: [
              {
                name: 'search_vendors',
                description: 'Search QuickBooks Online for existing vendors by trade category or company name. Returns contact info and outstanding balance.',
                parameters: {
                  trade: {
                    type: 'string',
                    description: 'Trade or service type to search for (e.g., "plumbing", "electrical", "HVAC", "janitorial", "landscaping")',
                    required: false,
                  },
                  name: {
                    type: 'string',
                    description: 'Vendor company name or partial name to search for',
                    required: false,
                  },
                },
              },
            ],
          },
        },
        {
          actionGroupName: 'Google_Maps_Lookup',
          description: 'Search Google Maps Places for vendors near a location. Use ONLY when QBO and the knowledge base have no suitable vendor.',
          actionGroupState: 'ENABLED',
          actionGroupExecutor: { lambda: this.mapsLambda.functionArn },
          functionSchema: {
            functions: [
              {
                name: 'search_nearby_vendors',
                description: 'Search Google Maps Places for local vendors and contractors by trade type and location.',
                parameters: {
                  trade: {
                    type: 'string',
                    description: 'Trade or service type to search for (e.g., "plumbing contractor", "electrician")',
                    required: true,
                  },
                  location: {
                    type: 'string',
                    description: 'Location to search near — can be a city, address, zip code, or facility/business name (e.g., "Seattle WA", "Amazon BFI9", "3230 International Pl DuPont WA"). Pass whatever location context the user provided; Google Maps will resolve it.',
                    required: true,
                  },
                },
              },
            ],
          },
        },
        {
          actionGroupName: 'WO_PO_Lookup',
          description: 'Look up work orders, purchase orders, and Amazon site assignments directly from the database',
          actionGroupState: 'ENABLED',
          actionGroupExecutor: { lambda: this.woPoLambda.functionArn },
          functionSchema: {
            functions: [
              {
                name: 'lookup_work_order',
                description: 'Look up a work order by its ID number. Returns status, description, site, assignment, dates, and full comment history.',
                parameters: {
                  work_order_id: {
                    type: 'string',
                    description: 'The work order ID number (e.g., "10046966057")',
                    required: true,
                  },
                },
              },
              {
                name: 'lookup_purchase_order',
                description: 'Look up a purchase order by its PO number. Returns status, supplier, ship-to, line items, amounts, and dates.',
                parameters: {
                  po_number: {
                    type: 'string',
                    description: 'The purchase order number (e.g., "2D-20023475")',
                    required: true,
                  },
                },
              },
              {
                name: 'lookup_site',
                description: 'Look up an Amazon facility site by its site code or list all sites in a state. Returns address, city, state, coordinates, and notes.',
                parameters: {
                  site_code: {
                    type: 'string',
                    description: 'The Amazon site code (e.g., "ABE2", "DFW6", "WWY1")',
                    required: false,
                  },
                  state: {
                    type: 'string',
                    description: 'Two-letter state abbreviation to list all sites in that state (e.g., "TX", "CA", "GA")',
                    required: false,
                  },
                },
              },
            ],
          },
        },
      ],
    });

    // ── Agent alias (stable ARN for invocations) ──────────────────────────────
    // Creating the alias also triggers agent preparation in CloudFormation.
    this.agentAlias = new bedrock.CfnAgentAlias(this, 'AgentAlias', {
      agentId: this.agent.attrAgentId,
      agentAliasName: 'live',
      description: 'Production alias — seahaven-assistant v8 (site lookups)',
    });
  }
}