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
+6 -5
View File
@@ -107,14 +107,15 @@ 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: