Use secrets manager for api key

Author: daixba

deployment/BedrockProxy.template

@@ -0,0 +1,279 @@
+Description: Bedrock Access Gateway - OpenAI-compatible RESTful APIs for Amazon Bedrock
+Parameters:
+  ApiKeySecretArn:
+    Type: String
+    AllowedPattern: ^arn:aws:secretsmanager:.*$
+    Description: The secret ARN in Secrets Manager used to store the API Key
+  DefaultModelId:
+    Type: String
+    Default: anthropic.claude-3-sonnet-20240229-v1:0
+    Description: The default model ID, please make sure the model ID is supported in the current region
+Resources:
+  VPCB9E5F0B4:
+    Type: AWS::EC2::VPC
+    Properties:
+      CidrBlock: 10.250.0.0/16
+      EnableDnsHostnames: true
+      EnableDnsSupport: true
+      InstanceTenancy: default
+      Tags:
+        - Key: Name
+          Value: BedrockProxy/VPC
+  VPCPublicSubnet1SubnetB4246D30:
+    Type: AWS::EC2::Subnet
+    Properties:
+      AvailabilityZone:
+        Fn::Select:
+          - 0
+          - Fn::GetAZs: ""
+      CidrBlock: 10.250.0.0/24
+      MapPublicIpOnLaunch: true
+      Tags:
+        - Key: aws-cdk:subnet-name
+          Value: Public
+        - Key: aws-cdk:subnet-type
+          Value: Public
+        - Key: Name
+          Value: BedrockProxy/VPC/PublicSubnet1
+      VpcId:
+        Ref: VPCB9E5F0B4
+  VPCPublicSubnet1RouteTableFEE4B781:
+    Type: AWS::EC2::RouteTable
+    Properties:
+      Tags:
+        - Key: Name
+          Value: BedrockProxy/VPC/PublicSubnet1
+      VpcId:
+        Ref: VPCB9E5F0B4
+  VPCPublicSubnet1RouteTableAssociation0B0896DC:
+    Type: AWS::EC2::SubnetRouteTableAssociation
+    Properties:
+      RouteTableId:
+        Ref: VPCPublicSubnet1RouteTableFEE4B781
+      SubnetId:
+        Ref: VPCPublicSubnet1SubnetB4246D30
+  VPCPublicSubnet1DefaultRoute91CEF279:
+    Type: AWS::EC2::Route
+    Properties:
+      DestinationCidrBlock: 0.0.0.0/0
+      GatewayId:
+        Ref: VPCIGWB7E252D3
+      RouteTableId:
+        Ref: VPCPublicSubnet1RouteTableFEE4B781
+    DependsOn:
+      - VPCVPCGW99B986DC
+  VPCPublicSubnet2Subnet74179F39:
+    Type: AWS::EC2::Subnet
+    Properties:
+      AvailabilityZone:
+        Fn::Select:
+          - 1
+          - Fn::GetAZs: ""
+      CidrBlock: 10.250.1.0/24
+      MapPublicIpOnLaunch: true
+      Tags:
+        - Key: aws-cdk:subnet-name
+          Value: Public
+        - Key: aws-cdk:subnet-type
+          Value: Public
+        - Key: Name
+          Value: BedrockProxy/VPC/PublicSubnet2
+      VpcId:
+        Ref: VPCB9E5F0B4
+  VPCPublicSubnet2RouteTable6F1A15F1:
+    Type: AWS::EC2::RouteTable
+    Properties:
+      Tags:
+        - Key: Name
+          Value: BedrockProxy/VPC/PublicSubnet2
+      VpcId:
+        Ref: VPCB9E5F0B4
+  VPCPublicSubnet2RouteTableAssociation5A808732:
+    Type: AWS::EC2::SubnetRouteTableAssociation
+    Properties:
+      RouteTableId:
+        Ref: VPCPublicSubnet2RouteTable6F1A15F1
+      SubnetId:
+        Ref: VPCPublicSubnet2Subnet74179F39
+  VPCPublicSubnet2DefaultRouteB7481BBA:
+    Type: AWS::EC2::Route
+    Properties:
+      DestinationCidrBlock: 0.0.0.0/0
+      GatewayId:
+        Ref: VPCIGWB7E252D3
+      RouteTableId:
+        Ref: VPCPublicSubnet2RouteTable6F1A15F1
+    DependsOn:
+      - VPCVPCGW99B986DC
+  VPCIGWB7E252D3:
+    Type: AWS::EC2::InternetGateway
+    Properties:
+      Tags:
+        - Key: Name
+          Value: BedrockProxy/VPC
+  VPCVPCGW99B986DC:
+    Type: AWS::EC2::VPCGatewayAttachment
+    Properties:
+      InternetGatewayId:
+        Ref: VPCIGWB7E252D3
+      VpcId:
+        Ref: VPCB9E5F0B4
+  ProxyApiHandlerServiceRoleBE71BFB1:
+    Type: AWS::IAM::Role
+    Properties:
+      AssumeRolePolicyDocument:
+        Statement:
+          - Action: sts:AssumeRole
+            Effect: Allow
+            Principal:
+              Service: lambda.amazonaws.com
+        Version: "2012-10-17"
+      ManagedPolicyArns:
+        - Fn::Join:
+            - ""
+            - - "arn:"
+              - Ref: AWS::Partition
+              - :iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
+  ProxyApiHandlerServiceRoleDefaultPolicy86681202:
+    Type: AWS::IAM::Policy
+    Properties:
+      PolicyDocument:
+        Statement:
+          - Action:
+              - bedrock:ListFoundationModels
+              - bedrock:ListInferenceProfiles
+            Effect: Allow
+            Resource: "*"
+          - Action:
+              - bedrock:InvokeModel
+              - bedrock:InvokeModelWithResponseStream
+            Effect: Allow
+            Resource:
+              - arn:aws:bedrock:*::foundation-model/*
+              - arn:aws:bedrock:*:*:inference-profile/*
+          - Action:
+              - secretsmanager:GetSecretValue
+              - secretsmanager:DescribeSecret
+            Effect: Allow
+            Resource:
+              Ref: ApiKeySecretArn
+        Version: "2012-10-17"
+      PolicyName: ProxyApiHandlerServiceRoleDefaultPolicy86681202
+      Roles:
+        - Ref: ProxyApiHandlerServiceRoleBE71BFB1
+  ProxyApiHandlerEC15A492:
+    Type: AWS::Lambda::Function
+    Properties:
+      Architectures:
+        - arm64
+      Code:
+        ImageUri:
+          Fn::Join:
+            - ""
+            - - 366590864501.dkr.ecr.
+              - Ref: AWS::Region
+              - "."
+              - Ref: AWS::URLSuffix
+              - /bedrock-proxy-api:latest
+      Description: Bedrock Proxy API Handler
+      Environment:
+        Variables:
+          DEBUG: "false"
+          API_KEY_SECRET_ARN:
+            Ref: ApiKeySecretArn
+          DEFAULT_MODEL:
+            Ref: DefaultModelId
+          DEFAULT_EMBEDDING_MODEL: cohere.embed-multilingual-v3
+          ENABLE_CROSS_REGION_INFERENCE: "true"
+      MemorySize: 1024
+      PackageType: Image
+      Role:
+        Fn::GetAtt:
+          - ProxyApiHandlerServiceRoleBE71BFB1
+          - Arn
+      Timeout: 600
+    DependsOn:
+      - ProxyApiHandlerServiceRoleDefaultPolicy86681202
+      - ProxyApiHandlerServiceRoleBE71BFB1
+  ProxyApiHandlerInvoke2UTWxhlfyqbT5FTn5jvgbLgjFfJwzswGk55DU1HYF6C33779:
+    Type: AWS::Lambda::Permission
+    Properties:
+      Action: lambda:InvokeFunction
+      FunctionName:
+        Fn::GetAtt:
+          - ProxyApiHandlerEC15A492
+          - Arn
+      Principal: elasticloadbalancing.amazonaws.com
+  ProxyALB87756780:
+    Type: AWS::ElasticLoadBalancingV2::LoadBalancer
+    Properties:
+      LoadBalancerAttributes:
+        - Key: deletion_protection.enabled
+          Value: "false"
+      Scheme: internet-facing
+      SecurityGroups:
+        - Fn::GetAtt:
+            - ProxyALBSecurityGroup0D6CA3DA
+            - GroupId
+      Subnets:
+        - Ref: VPCPublicSubnet1SubnetB4246D30
+        - Ref: VPCPublicSubnet2Subnet74179F39
+      Type: application
+    DependsOn:
+      - VPCPublicSubnet1DefaultRoute91CEF279
+      - VPCPublicSubnet1RouteTableAssociation0B0896DC
+      - VPCPublicSubnet2DefaultRouteB7481BBA
+      - VPCPublicSubnet2RouteTableAssociation5A808732
+  ProxyALBSecurityGroup0D6CA3DA:
+    Type: AWS::EC2::SecurityGroup
+    Properties:
+      GroupDescription: Automatically created Security Group for ELB BedrockProxyALB1CE4CAD1
+      SecurityGroupEgress:
+        - CidrIp: 255.255.255.255/32
+          Description: Disallow all traffic
+          FromPort: 252
+          IpProtocol: icmp
+          ToPort: 86
+      SecurityGroupIngress:
+        - CidrIp: 0.0.0.0/0
+          Description: Allow from anyone on port 80
+          FromPort: 80
+          IpProtocol: tcp
+          ToPort: 80
+      VpcId:
+        Ref: VPCB9E5F0B4
+  ProxyALBListener933E9515:
+    Type: AWS::ElasticLoadBalancingV2::Listener
+    Properties:
+      DefaultActions:
+        - TargetGroupArn:
+            Ref: ProxyALBListenerTargetsGroup187739FA
+          Type: forward
+      LoadBalancerArn:
+        Ref: ProxyALB87756780
+      Port: 80
+      Protocol: HTTP
+  ProxyALBListenerTargetsGroup187739FA:
+    Type: AWS::ElasticLoadBalancingV2::TargetGroup
+    Properties:
+      HealthCheckEnabled: false
+      TargetType: lambda
+      Targets:
+        - Id:
+            Fn::GetAtt:
+              - ProxyApiHandlerEC15A492
+              - Arn
+    DependsOn:
+      - ProxyApiHandlerInvoke2UTWxhlfyqbT5FTn5jvgbLgjFfJwzswGk55DU1HYF6C33779
+Outputs:
+  APIBaseUrl:
+    Description: Proxy API Base URL (OPENAI_API_BASE)
+    Value:
+      Fn::Join:
+        - ""
+        - - http://
+          - Fn::GetAtt:
+              - ProxyALB87756780
+              - DNSName
+          - /api/v1
+

deployment/BedrockProxyFargate.template

@@ -1,10 +1,13 @@
 Description: Bedrock Access Gateway - OpenAI-compatible RESTful APIs for Amazon Bedrock
-Transform: AWS::LanguageExtensions
 Parameters:
   ApiKeySecretArn:
     Type: String
     AllowedPattern: ^arn:aws:secretsmanager:.*$
     Description: The secret ARN in Secrets Manager used to store the API Key
+  DefaultModelId:
+    Type: String
+    Default: anthropic.claude-3-sonnet-20240229-v1:0
+    Description: The default model ID, please make sure the model ID is supported in the current region
 Resources:
   VPCB9E5F0B4:
     Type: AWS::EC2::VPC
@@ -214,11 +217,7 @@ Resources:
               Value: "false"
             - Name: DEFAULT_MODEL
               Value:
-                Fn::FindInMap:
-                  - ProxyRegionTable03E5BEB3
-                  - Ref: AWS::Region
-                  - model
-                  - DefaultValue: anthropic.claude-3-sonnet-20240229-v1:0
+                Ref: DefaultModelId
             - Name: DEFAULT_EMBEDDING_MODEL
               Value: cohere.embed-multilingual-v3
             - Name: ENABLE_CROSS_REGION_INFERENCE
@@ -407,10 +406,6 @@ Resources:
       TargetType: ip
       VpcId:
         Ref: VPCB9E5F0B4
-Mappings:
-  ProxyRegionTable03E5BEB3:
-    us-east-1:
-      model: anthropic.claude-3-sonnet-20240229-v1:0
 Outputs:
   APIBaseUrl:
     Description: Proxy API Base URL (OPENAI_API_BASE)

src/api/auth.py

@@ -1,22 +1,41 @@
+import json
 import os
 from typing import Annotated
 
 import boto3
+from botocore.exceptions import ClientError
 from fastapi import Depends, HTTPException, status
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 
 from api.setting import DEFAULT_API_KEYS
 
 api_key_param = os.environ.get("API_KEY_PARAM_NAME")
+api_key_secret_arn = os.environ.get("API_KEY_SECRET_ARN")
 api_key_env = os.environ.get("API_KEY")
 if api_key_param:
+    # For backward compatibility.
+    # Please now use secrets manager instead.
     ssm = boto3.client("ssm")
     api_key = ssm.get_parameter(Name=api_key_param, WithDecryption=True)["Parameter"][
         "Value"
     ]
+elif api_key_secret_arn:
+    sm = boto3.client("secretsmanager")
+    try:
+        response = sm.get_secret_value(SecretId=api_key_secret_arn)
+        if "SecretString" in response:
+            secret = json.loads(response["SecretString"])
+            api_key = secret["api_key"]
+    except ClientError as e:
+        raise RuntimeError(
+            "Unable to retrieve API KEY, please ensure the secret ARN is correct"
+        )
+    except KeyError as e:
+        raise RuntimeError('Please ensure the secret contains a "api_key" field')
 elif api_key_env:
     api_key = api_key_env
 else:
+    # For local use only.
     api_key = DEFAULT_API_KEYS
 
 security = HTTPBearer()

src/api/models/bedrock.py

@@ -105,7 +105,7 @@ def list_bedrock_models() -> dict:
             status = model['modelLifecycle'].get('status', 'ACTIVE')
 
             # currently, use this to filter out rerank models and legacy models
-            if not stream_supported or status != "ACTIVE":
+            if not stream_supported or status not in ["ACTIVE", "LEGACY"]:
                 continue
 
             inference_types = model.get('inferenceTypesSupported', [])