mirror of
https://github.com/open-webui/pipelines
synced 2025-05-10 15:40:55 +00:00
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.
This commit is contained in:
parent
c1bbbe1165
commit
ecc44ebd1e
@ -100,14 +100,18 @@ class Pipeline:
|
||||
|
||||
def get_models(self):
|
||||
try:
|
||||
response = self.bedrock.list_foundation_models(byProvider='Anthropic', byInferenceType='ON_DEMAND')
|
||||
return [
|
||||
{
|
||||
"id": model["modelId"],
|
||||
"name": model["modelName"],
|
||||
}
|
||||
for model in response["modelSummaries"]
|
||||
]
|
||||
res = []
|
||||
response = self.bedrock.list_foundation_models(byProvider='Anthropic')
|
||||
for model in response['modelSummaries']:
|
||||
inference_types = model.get('inferenceTypesSupported', [])
|
||||
if "ON_DEMAND" in inference_types:
|
||||
res.append({'id': model['modelId'], 'name': model['modelName']})
|
||||
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]:
|
||||
|
Loading…
Reference in New Issue
Block a user