Enhance get_models method to include models with INFERENCE_PROFILE type

- Updated the get_models method to fetch models that support both ON_DEMAND and INFERENCE_PROFILE inference types.
- Added a helper method getInferenceProfileId to retrieve the inference profile ID for models with INFERENCE_PROFILE type.
- This change ensures that models with different inference types are correctly listed and available for use.

examples/pipelines/providers/aws_bedrock_claude_pipeline.py

@@ -100,14 +100,18 @@ class Pipeline:
 
     def get_models(self):
         try:
-            response = self.bedrock.list_foundation_models(byProvider='Anthropic', byInferenceType='ON_DEMAND')
+            res = []
-            return [
+            response = self.bedrock.list_foundation_models(byProvider='Anthropic')
-                {
+            for model in response['modelSummaries']:
-                    "id": model["modelId"],
+                inference_types = model.get('inferenceTypesSupported', [])
-                    "name": model["modelName"],
+                if "ON_DEMAND" in inference_types:
-                }
+                    res.append({'id': model['modelId'], 'name': model['modelName']})
-                for model in response["modelSummaries"]
+                elif "INFERENCE_PROFILE" in inference_types:
-            ]
+                    inferenceProfileId = self.getInferenceProfileId(model['modelArn'])
+                    if inferenceProfileId:
+                        res.append({'id': inferenceProfileId, 'name': model['modelName']})
+
+            return res
         except Exception as e:
             print(f"Error: {e}")
             return [
@@ -117,6 +121,14 @@ class Pipeline:
                 },
             ]
 
+    def getInferenceProfileId(self, modelArn: str) -> str:
+        response = self.bedrock.list_inference_profiles()
+        for profile in response.get('inferenceProfileSummaries', []):
+            for model in profile.get('models', []):
+                if model.get('modelArn') == modelArn:
+                    return profile['inferenceProfileId']
+        return None
+
     def pipe(
         self, user_message: str, model_id: str, messages: List[dict], body: dict
     ) -> Union[str, Generator, Iterator]: