feat(ai): add cross-region inference support for conversation routes (#2951)

Author: esslam-ashour

.changeset/yellow-corners-like.md

@@ -0,0 +1,11 @@
+---
+'@aws-amplify/ai-constructs': minor
+'@aws-amplify/backend-ai': minor
+---
+
+feat(ai): add cross-region inference support for AI conversation routes
+
+- Add `AiModelArnGeneratorConstruct` + `AiModelPropsResolver` for region-aware model/inference-profile resolution.
+- Conversation handler supports `crossRegionInference` and uses generated ARNs for IAM policies.
+- Runtime (`BedrockConverseAdapter`) resolves foundation model vs inference profile ID using `AiModelPropsResolver`.
+- Add `models.json` manifest and tests; update public exports.

package-lock.json

@@ -14,6 +14,23 @@ import * as jsonSchemaToTypeScript from 'json-schema-to-ts';
 import { ResourceProvider } from '@aws-amplify/plugin-types';
 import { RetentionDays } from 'aws-cdk-lib/aws-logs';
 
+declare namespace __export__ai_model {
+    export {
+        AiModelArnGeneratorConstruct,
+        AiModelPropsResolver,
+        AiModelConfig
+    }
+}
+export { __export__ai_model }
+
+declare namespace __export__ai_model__runtime {
+    export {
+        AiModelPropsResolver,
+        AiModelConfig
+    }
+}
+export { __export__ai_model__runtime }
+
 declare namespace __export__conversation {
     export {
         ConversationHandlerFunction,
@@ -38,6 +55,33 @@ declare namespace __export__conversation__runtime {
 }
 export { __export__conversation__runtime }
 
+// @public
+class AiModelArnGeneratorConstruct extends Construct {
+    constructor(scope: Construct, id?: string);
+    generateArns(modelConfig: AiModelConfig): string[];
+}
+
+// @public
+type AiModelConfig = {
+    modelId: string;
+    region: string;
+    crossRegionInference: boolean;
+};
+
+// @public
+class AiModelPropsResolver {
+    constructor();
+    getFoundationModelId(inferenceProfileId: string): string;
+    getGeography(region: string): string;
+    getInferenceProfileId(foundationModelId: string, geography: string): string;
+    getSupportedSourceRegions(modelId: string, geography: string): string[];
+    isKnownInferenceProfile(modelId: string): boolean;
+    requiresCri(modelId: string, region: string): boolean;
+    resolveModelId(modelConfig: AiModelConfig): string;
+    supportsCri(modelId: string, region: string): boolean;
+    validateModelId(modelId: string): void;
+}
+
 // @public
 class ConversationHandlerFunction extends Construct implements ResourceProvider<FunctionResources> {
     constructor(scope: Construct, id: string, props: ConversationHandlerFunctionProps);
@@ -52,6 +96,7 @@ type ConversationHandlerFunctionProps = {
     entry?: string;
     models: Array<{
         modelId: string;
+        crossRegionInference?: boolean;
         region?: string;
     }>;
     memoryMB?: number;
@@ -76,6 +121,7 @@ type ConversationTurnEvent = {
     graphqlApiEndpoint: string;
     modelConfiguration: {
         modelId: string;
+        crossRegionInference?: boolean;
         systemPrompt: string;
         region?: string;
         inferenceConfiguration?: {

packages/ai-constructs/API.md

@@ -17,6 +17,14 @@
     "./conversation/runtime": {
       "types": "./lib/conversation/runtime/index.d.ts",
       "default": "./lib/conversation/runtime/index.js"
+    },
+    "./ai-model": {
+      "types": "./lib/ai-model/index.d.ts",
+      "default": "./lib/ai-model/index.js"
+    },
+    "./ai-model/runtime": {
+      "types": "./lib/ai-model/runtime/index.d.ts",
+      "default": "./lib/ai-model/runtime/index.js"
     }
   },
   "main": "lib/index.js",
@@ -36,7 +44,9 @@
   },
   "devDependencies": {
     "@aws-amplify/backend-output-storage": "^1.3.1",
+    "@types/aws-lambda": "^8.10.119",
     "@types/lodash.transform": "^4.6.9",
+    "aws-lambda": "^1.0.7",
     "lodash.transform": "^4.6.0",
     "typescript": "^5.0.0"
   },

packages/ai-constructs/package.json

@@ -0,0 +1,104 @@
+import { describe, it } from 'node:test';
+import assert from 'node:assert';
+import { App, Stack } from 'aws-cdk-lib';
+import { Template } from 'aws-cdk-lib/assertions';
+import { AiModelArnGeneratorConstruct } from './ai_model_arn_generator_construct';
+import { TEST_MODEL_IDS } from './test-assets';
+
+void describe('AiModelArnGeneratorConstruct', () => {
+  const { FOUNDATION } = TEST_MODEL_IDS;
+
+  const createStack = () => {
+    const app = new App();
+    return new Stack(app, 'TestStack');
+  };
+
+  void describe('constructor', () => {
+    void it('creates construct with default id', () => {
+      const stack = createStack();
+      const construct = new AiModelArnGeneratorConstruct(stack);
+
+      assert.ok(construct);
+      assert.strictEqual(construct.node.id, 'AmplifyAiModelArnGenerator');
+    });
+
+    void it('creates construct with custom id', () => {
+      const stack = createStack();
+      const construct = new AiModelArnGeneratorConstruct(stack, 'CustomId');
+
+      assert.ok(construct);
+      assert.strictEqual(construct.node.id, 'CustomId');
+    });
+
+    void it('creates Lambda function', () => {
+      const stack = createStack();
+      new AiModelArnGeneratorConstruct(stack);
+
+      const template = Template.fromStack(stack);
+
+      // Should create Lambda function
+      template.hasResourceProperties('AWS::Lambda::Function', {
+        Runtime: 'nodejs20.x',
+        Timeout: 30,
+        MemorySize: 128,
+        Description:
+          'Resolve Amazon Bedrock AI model ARNs based on region and cross-region inference settings',
+      });
+    });
+  });
+
+  void describe('generateArns', () => {
+    void it('returns array of ARN tokens', () => {
+      const stack = createStack();
+      const construct = new AiModelArnGeneratorConstruct(stack);
+
+      const modelConfig = {
+        modelId: FOUNDATION,
+        region: 'us-east-1',
+        crossRegionInference: false,
+      };
+
+      const arns = construct.generateArns(modelConfig);
+
+      assert.ok(Array.isArray(arns));
+      assert.strictEqual(arns.length, 1);
+    });
+
+    void it('returns array when generating ARNs', () => {
+      const stack = createStack();
+      const construct = new AiModelArnGeneratorConstruct(stack);
+
+      const modelConfig = {
+        modelId: FOUNDATION,
+        region: 'us-east-1',
+        crossRegionInference: false,
+      };
+
+      const arns = construct.generateArns(modelConfig);
+
+      assert.ok(Array.isArray(arns));
+      assert.strictEqual(arns.length, 1);
+    });
+
+    void it('returns different arrays for different calls', () => {
+      const stack = createStack();
+      const construct = new AiModelArnGeneratorConstruct(stack);
+
+      const arns1 = construct.generateArns({
+        modelId: FOUNDATION,
+        region: 'us-east-1',
+        crossRegionInference: false,
+      });
+
+      const arns2 = construct.generateArns({
+        modelId: FOUNDATION,
+        region: 'us-west-2',
+        crossRegionInference: true,
+      });
+
+      assert.ok(Array.isArray(arns1));
+      assert.ok(Array.isArray(arns2));
+      assert.notStrictEqual(arns1, arns2);
+    });
+  });
+});

packages/ai-constructs/src/ai-model/ai_model_arn_generator_construct.test.ts

@@ -0,0 +1,75 @@
+import { Construct } from 'constructs';
+import { NodejsFunction } from 'aws-cdk-lib/aws-lambda-nodejs';
+import { Runtime as LambdaRuntime } from 'aws-cdk-lib/aws-lambda';
+import { CustomResource, Duration, Fn } from 'aws-cdk-lib';
+import { Provider } from 'aws-cdk-lib/custom-resources';
+import { RetentionDays } from 'aws-cdk-lib/aws-logs';
+import { resolve } from 'path';
+import { AiModelConfig } from './ai_model_types';
+
+/**
+ * Custom resource type identifier.
+ */
+const AI_MODEL_ARN_GENERATOR_RESOURCE_TYPE =
+  'Custom::AmplifyAiModelArnGenerator';
+
+const DEFAULT_LAMBDA_TIMEOUT_SECONDS = 30;
+const DEFAULT_LAMBDA_MEMORY_SIZE = 128;
+const DEFAULT_LOG_RETENTION_DAYS = RetentionDays.TWO_WEEKS;
+
+/**
+ * CDK construct that generates AI model ARNs using a custom resource.
+ * Handles both foundation models and inference profiles based on configuration.
+ */
+export class AiModelArnGeneratorConstruct extends Construct {
+  private readonly provider: Provider;
+  private resourceCounter = 0;
+
+  /**
+   * Creates a new AI model ARN generator construct.
+   */
+  constructor(scope: Construct, id = 'AmplifyAiModelArnGenerator') {
+    super(scope, id);
+
+    const entryPath = resolve(__dirname, 'lambda', 'ai_model_arn_generator.js');
+
+    const onEvent = new NodejsFunction(this, 'AiModelArnGeneratorLambda', {
+      runtime: LambdaRuntime.NODEJS_20_X,
+      timeout: Duration.seconds(DEFAULT_LAMBDA_TIMEOUT_SECONDS),
+      entry: entryPath,
+      handler: 'handler',
+      description:
+        'Resolve Amazon Bedrock AI model ARNs based on region and cross-region inference settings',
+      memorySize: DEFAULT_LAMBDA_MEMORY_SIZE,
+      logRetention: DEFAULT_LOG_RETENTION_DAYS,
+      bundling: {
+        // Include deps to avoid relying on Lambda layer versions.
+        externalModules: [],
+      },
+    });
+
+    this.provider = new Provider(this, 'AiModelArnGeneratorProvider', {
+      onEventHandler: onEvent,
+      logRetention: DEFAULT_LOG_RETENTION_DAYS,
+    });
+  }
+
+  /**
+   * Generates model ARNs based on the provided configuration.
+   * Returns CloudFormation tokens that resolve to an array of ARN strings.
+   */
+  public generateArns(modelConfig: AiModelConfig): string[] {
+    this.resourceCounter += 1;
+    const resourceId = `AiModelArnsResource${this.resourceCounter}`;
+
+    const resource = new CustomResource(this, resourceId, {
+      serviceToken: this.provider.serviceToken,
+      resourceType: AI_MODEL_ARN_GENERATOR_RESOURCE_TYPE,
+      properties: {
+        modelConfig,
+      },
+    });
+
+    return Fn.split(',', resource.getAtt('modelArns').toString());
+  }
+}

packages/ai-constructs/src/ai-model/ai_model_arn_generator_construct.ts

@@ -0,0 +1,32 @@
+/** Mapping of geography names to their regions. */
+export type RegionGroups = Record<string, string[]>;
+
+/** Configuration for a model within a specific geography. */
+export type AiModelMapping = {
+  /** Inference profile ID for cross-region inference, null if not available. */
+  inferenceProfileId: string | null;
+  /** Regions that can source requests for this model using cross-region inference. */
+  supportsCriFrom: string[];
+  /** Regions where cross-region inference is required. */
+  requiresCriIn: string[];
+};
+
+/** Model entry containing geography-specific configurations. */
+export type AiModelsEntry = {
+  geographies: Record<string, AiModelMapping>;
+};
+
+/** Root structure of the models JSON data. */
+export type AiModelsManifest = {
+  schemaVersion: 1;
+  lastUpdated: string;
+  regionGroups: RegionGroups;
+  models: Record<string, AiModelsEntry>;
+};
+
+/** Configuration for AI model ARN and ID generation. */
+export type AiModelConfig = {
+  modelId: string;
+  region: string;
+  crossRegionInference: boolean;
+};

packages/ai-constructs/src/ai-model/ai_model_types.ts

@@ -0,0 +1,3 @@
+export { AiModelArnGeneratorConstruct } from './ai_model_arn_generator_construct';
+export { AiModelPropsResolver } from './runtime/ai_model_props_resolver';
+export type { AiModelConfig } from './ai_model_types';

packages/ai-constructs/src/ai-model/index.ts

@@ -0,0 +1,5 @@
+{
+  "rules": {
+    "no-console": "off"
+  }
+}