AVB/bedrock-access-gateway
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add pagination to list_inference_profiles calls (#173)

Browse Source 
Co-authored-by: Rizvi Rahim <rizvi@rizvir.com>
This commit is contained in:
RizviR
2025-08-13 10:26:34 +08:00
committed by GitHub
co-authored by Rizvi Rahim
parent 0cce2edab0
commit a2110ff648
1 changed files with 21 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/api/models/bedrock.py
+21 -20
View File
@@ -107,31 +107,32 @@ def list_bedrock_models() -> dict:
if ENABLE_CROSS_REGION_INFERENCE: if ENABLE_CROSS_REGION_INFERENCE:
# List system defined inference profile IDs # List system defined inference profile IDs
response = bedrock_client.list_inference_profiles(maxResults=1000, typeEquals="SYSTEM_DEFINED") paginator = bedrock_client.get_paginator('list_inference_profiles')
profile_list = [p["inferenceProfileId"] for p in response["inferenceProfileSummaries"]] for page in paginator.paginate(maxResults=1000, typeEquals="SYSTEM_DEFINED"):
profile_list.extend([p["inferenceProfileId"] for p in page["inferenceProfileSummaries"]])
if ENABLE_APPLICATION_INFERENCE_PROFILES: if ENABLE_APPLICATION_INFERENCE_PROFILES:
# List application defined inference profile IDs and create mapping # List application defined inference profile IDs and create mapping
response = bedrock_client.list_inference_profiles(maxResults=1000, typeEquals="APPLICATION") paginator = bedrock_client.get_paginator('list_inference_profiles')
for page in paginator.paginate(maxResults=1000, typeEquals="APPLICATION"):
for profile in response["inferenceProfileSummaries"]: for profile in page["inferenceProfileSummaries"]:
try: try:
profile_arn = profile.get("inferenceProfileArn") profile_arn = profile.get("inferenceProfileArn")
if not profile_arn: if not profile_arn:
continue
# Process all models in the profile
models = profile.get("models", [])
for model in models:
model_arn = model.get("modelArn", "")
if model_arn:
model_id = model_arn.split('/')[-1] if '/' in model_arn else model_arn
if model_id:
app_profile_dict[model_id] = profile_arn
except Exception as e:
logger.warning(f"Error processing application profile: {e}")
continue continue
# Process all models in the profile
models = profile.get("models", [])
for model in models:
model_arn = model.get("modelArn", "")
if model_arn:
model_id = model_arn.split('/')[-1] if '/' in model_arn else model_arn
if model_id:
app_profile_dict[model_id] = profile_arn
except Exception as e:
logger.warning(f"Error processing application profile: {e}")
continue
# List foundation models, only cares about text outputs here. # List foundation models, only cares about text outputs here.
response = bedrock_client.list_foundation_models(byOutputModality="TEXT") response = bedrock_client.list_foundation_models(byOutputModality="TEXT")