Support imported models

Signed-off-by: Sean Smith <[email protected]>

Author: Sean Smith

deployment/BedrockProxy.template renamed to deployment/BedrockProxy.yaml

@@ -5,6 +5,13 @@ Parameters:
     Type: String
     Default: ""
     Description: The parameter name in System Manager used to store the API Key, leave blank to use a default key
+  EnableImportedModels:
+    Type: String
+    Default: false
+    AllowedValues:
+      - true
+      - false
+    Description: If enabled, models imported into Bedrock will be available to use.
 Resources:
   VPCB9E5F0B4:
     Type: AWS::EC2::VPC
@@ -197,6 +204,7 @@ Resources:
               - DefaultValue: anthropic.claude-3-sonnet-20240229-v1:0
           DEFAULT_EMBEDDING_MODEL: cohere.embed-multilingual-v3
           ENABLE_CROSS_REGION_INFERENCE: "true"
+          ENABLE_IMPORTED_MODELS: !Ref EnableImportedModels
       MemorySize: 1024
       PackageType: Image
       Role:

deployment/BedrockProxyFargate.template renamed to deployment/BedrockProxyFargate.yaml

@@ -5,6 +5,13 @@ Parameters:
     Type: String
     Default: ""
     Description: The parameter name in System Manager used to store the API Key, leave blank to use a default key
+  EnableImportedModels:
+    Type: String
+    Default: false
+    AllowedValues:
+      - true
+      - false
+    Description: If enabled, models imported into Bedrock will be available to use.
 Resources:
   VPCB9E5F0B4:
     Type: AWS::EC2::VPC
@@ -237,6 +244,8 @@ Resources:
               Value: cohere.embed-multilingual-v3
             - Name: ENABLE_CROSS_REGION_INFERENCE
               Value: "true"
+            - Name: ENABLE_IMPORTED_MODELS
+              Value: !Ref EnableImportedModels
           Essential: true
           Image:
             Fn::Join:

src/api/models/bedrock.py

@@ -38,7 +38,7 @@
     Embedding,
 
 )
-from api.setting import DEBUG, AWS_REGION, ENABLE_CROSS_REGION_INFERENCE, DEFAULT_MODEL
+from api.setting import DEBUG, AWS_REGION, ENABLE_CROSS_REGION_INFERENCE, DEFAULT_MODEL, ENABLE_IMPORTED_MODELS
 
 logger = logging.getLogger(__name__)
 
@@ -99,6 +99,17 @@ def list_bedrock_models() -> dict:
             byOutputModality='TEXT'
         )
 
+        # Add imported models to the list if ENABLE_IMPORTED_MODELS is true
+        if ENABLE_IMPORTED_MODELS:
+            response_imported = bedrock_client.list_imported_models()
+
+            # Add imported models to the default model list
+            for model in response_imported['modelSummaries']:
+                model_id = model.get('modelArn')
+                model_list[model_id] = {
+                    'modalities': ["TEXT"]
+                }
+
         for model in response['modelSummaries']:
             model_id = model.get('modelId', 'N/A')
             stream_supported = model.get('responseStreamingSupported', True)

src/api/setting.py

@@ -20,3 +20,4 @@
     "DEFAULT_EMBEDDING_MODEL", "cohere.embed-multilingual-v3"
 )
 ENABLE_CROSS_REGION_INFERENCE = os.environ.get("ENABLE_CROSS_REGION_INFERENCE", "true").lower() != "false"
+ENABLE_IMPORTED_MODELS = os.environ.get("ENABLE_IMPORTED_MODELS", "true").lower() != "false"