feature: inference profiles for bedrock #1118
Conversation
|
DescriptionSummary By MatterAI
🔄 What ChangedThis pull request introduces support for AWS Bedrock inference profiles. The core change involves modifying the 🔍 Impact of the ChangeThis feature allows users to leverage AWS Bedrock's inference profiles, providing a more flexible and potentially managed way to access foundation models without directly specifying the model ID. By caching the inference profile lookup for up to a day, the change significantly reduces latency and API calls to AWS Bedrock for repeated requests using the same inference profile. This enhances the overall performance and usability for Bedrock users. 📁 Total Files Changed8 files were changed in this pull request. 🧪 Test Added
🔒Security VulnerabilitiesNo security vulnerabilities were detected in the provided code patch. MotivationThis feature was motivated by the need to support AWS Bedrock's inference profiles, which provide a more flexible and managed way to interact with foundation models. This allows users to specify an inference profile ARN instead of a direct model identifier, enhancing integration with AWS Bedrock's advanced features and potentially simplifying model management for users. Type of Change
How Has This Been Tested?
Screenshots (if applicable)N/A Checklist
Related Issues. Tip Quality Recommendations
Sequence DiagramsequenceDiagram
participant Client
participant HandlerUtils
participant BedrockAPI
participant BedrockUtils
participant Cache
participant AWSBedrockAPI
Client->>HandlerUtils: POST /v1/chat/completions (model: inferenceProfileARN)
HandlerUtils: tryPost(fn, c, gatewayRequestURL, params)
HandlerUtils->>BedrockAPI: BedrockAPIConfig.getBaseURL({c, providerOptions, fn, gatewayRequestURL, params})
BedrockAPI: getBaseURL is now async
BedrockAPI->>BedrockAPI: Decode model from params (e.g., model = 'arn:aws:...')
alt model is an ARN and includes 'arn:aws'
BedrockAPI->>BedrockUtils: getFoundationModelFromInferenceProfile(c, model, providerOptions)
BedrockUtils->>Cache: getFromCacheByKey(env(c), cacheKey)
alt Cache Hit
Cache-->>BedrockUtils: cachedFoundationModel
BedrockUtils-->>BedrockAPI: cachedFoundationModel
else Cache Miss
BedrockUtils->>BedrockUtils: getInferenceProfile(inferenceProfileIdentifier, awsRegion, awsAccessKeyId, ...)
BedrockUtils->>BedrockUtils: generateAWSHeaders(..., url, 'GET', 'bedrock', ...)
BedrockUtils->>AWSBedrockAPI: GET /inference-profiles/{identifier} (with AWS headers)
AWSBedrockAPI-->>BedrockUtils: BedrockInferenceProfile JSON
BedrockUtils: Extract foundationModel from inferenceProfile.models[0].modelArn
BedrockUtils->>Cache: putInCacheWithValue(env(c), cacheKey, foundationModel, 86400)
BedrockUtils-->>BedrockAPI: foundationModel
end
BedrockAPI: Set params.foundationModel = foundationModel
end
BedrockAPI-->>HandlerUtils: baseURL
HandlerUtils->>BedrockAPI: Continue with API call using baseURL
BedrockAPI->>BedrockUtils: BedrockConfig (uses params.foundationModel if available)
Note over BedrockAPI,BedrockUtils: Other handlers (getBatchOutput, retrieveFileContent) also await getBaseURL
|
![matter-code-review[bot]](https://avatars.githubusercontent.com/in/474549?s=60&v=4){kind=small}
There was a problem hiding this comment.
This PR adds support for AWS Bedrock inference profiles, allowing the system to work with inference profile ARNs by extracting the underlying foundation model. The implementation is generally good with proper caching and error handling, but I've identified a few improvements that could enhance the code quality and reliability.
src/providers/bedrock/utils.ts
Outdated
| try { | ||
| const getFromCacheByKey = c.get('getFromCacheByKey'); | ||
| const putInCacheWithValue = c.get('putInCacheWithValue'); | ||
| const cacheKey = `bedrock-inference-profile-${inferenceProfileIdentifier}`; | ||
| const cachedFoundationModel = getFromCacheByKey | ||
| ? await getFromCacheByKey(env(c), cacheKey) | ||
| : null; | ||
| if (cachedFoundationModel) { | ||
| //update ttl, dont't await the result | ||
| putInCacheWithValue(env(c), cacheKey, cachedFoundationModel, 56400); | ||
| return cachedFoundationModel; | ||
| } | ||
|
|
||
| const inferenceProfile = await getInferenceProfile( | ||
| inferenceProfileIdentifier || '', | ||
| providerOptions.awsRegion || '', | ||
| providerOptions.awsAccessKeyId || '', | ||
| providerOptions.awsSecretAccessKey || '', | ||
| providerOptions.awsSessionToken || '' | ||
| ); | ||
|
|
||
| // modelArn is always like arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-v2:1 | ||
| const foundationModel = inferenceProfile?.models?.[0]?.modelArn | ||
| ?.split('/') | ||
| ?.pop(); | ||
| putInCacheWithValue(env(c), cacheKey, foundationModel, 56400); | ||
| return foundationModel; | ||
| } catch (error) { | ||
| return null; | ||
| } |
There was a problem hiding this comment.
🛠️ Code Refactor
Issue: The error handling in getFoundationModelFromInferenceProfile silently returns null for any error, which could hide important issues.
Fix: Add more specific error handling and logging to help with debugging.
Impact: Improves troubleshooting and error visibility when inference profile resolution fails.
| try { | |
| const getFromCacheByKey = c.get('getFromCacheByKey'); | |
| const putInCacheWithValue = c.get('putInCacheWithValue'); | |
| const cacheKey = `bedrock-inference-profile-${inferenceProfileIdentifier}`; | |
| const cachedFoundationModel = getFromCacheByKey | |
| ? await getFromCacheByKey(env(c), cacheKey) | |
| : null; | |
| if (cachedFoundationModel) { | |
| //update ttl, dont't await the result | |
| putInCacheWithValue(env(c), cacheKey, cachedFoundationModel, 56400); | |
| return cachedFoundationModel; | |
| } | |
| const inferenceProfile = await getInferenceProfile( | |
| inferenceProfileIdentifier || '', | |
| providerOptions.awsRegion || '', | |
| providerOptions.awsAccessKeyId || '', | |
| providerOptions.awsSecretAccessKey || '', | |
| providerOptions.awsSessionToken || '' | |
| ); | |
| // modelArn is always like arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-v2:1 | |
| const foundationModel = inferenceProfile?.models?.[0]?.modelArn | |
| ?.split('/') | |
| ?.pop(); | |
| putInCacheWithValue(env(c), cacheKey, foundationModel, 56400); | |
| return foundationModel; | |
| } catch (error) { | |
| return null; | |
| } | |
| try { | |
| const getFromCacheByKey = c.get('getFromCacheByKey'); | |
| const putInCacheWithValue = c.get('putInCacheWithValue'); | |
| const cacheKey = `bedrock-inference-profile-${inferenceProfileIdentifier}`; | |
| const cachedFoundationModel = getFromCacheByKey | |
| ? await getFromCacheByKey(env(c), cacheKey) | |
| : null; | |
| if (cachedFoundationModel) { | |
| //update ttl, dont't await the result | |
| putInCacheWithValue(env(c), cacheKey, cachedFoundationModel, 56400); | |
| return cachedFoundationModel; | |
| } | |
| const inferenceProfile = await getInferenceProfile( | |
| inferenceProfileIdentifier || '', | |
| providerOptions.awsRegion || '', | |
| providerOptions.awsAccessKeyId || '', | |
| providerOptions.awsSecretAccessKey || '', | |
| providerOptions.awsSessionToken || '' | |
| ); | |
| // modelArn is always like arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-v2:1 | |
| const foundationModel = inferenceProfile?.models?.[0]?.modelArn | |
| ?.split('/') | |
| ?.pop(); | |
| if (!foundationModel) { | |
| console.warn(`No foundation model found in inference profile: ${inferenceProfileIdentifier}`); | |
| return null; | |
| } | |
| putInCacheWithValue(env(c), cacheKey, foundationModel, 56400); | |
| return foundationModel; | |
| } catch (error) { | |
| console.error(`Error resolving foundation model from inference profile ${inferenceProfileIdentifier}:`, error); | |
| return null; | |
| } |
| const foundationModel = model.includes('foundation-model/') | ||
| ? model.split('/').pop() | ||
| : await getFoundationModelFromInferenceProfile( | ||
| c, | ||
| model, | ||
| providerOptions | ||
| ); | ||
| if (foundationModel) { | ||
| params.foundationModel = foundationModel; | ||
| } |
There was a problem hiding this comment.
🛠️ Code Refactor
Issue: The code doesn't handle the case where foundationModel extraction fails but still attempts to use it.
Fix: Add a check to ensure foundationModel is defined before setting it in params.
Impact: Prevents potential undefined values from being used in the model parameter.
| const foundationModel = model.includes('foundation-model/') | |
| ? model.split('/').pop() | |
| : await getFoundationModelFromInferenceProfile( | |
| c, | |
| model, | |
| providerOptions | |
| ); | |
| if (foundationModel) { | |
| params.foundationModel = foundationModel; | |
| } | |
| const foundationModel = model.includes('foundation-model/') | |
| ? model.split('/').pop() | |
| : await getFoundationModelFromInferenceProfile( | |
| c, | |
| model, | |
| providerOptions | |
| ); | |
| if (foundationModel && foundationModel.length > 0) { | |
| params.foundationModel = foundationModel; | |
| } |
| export const getInferenceProfile = async ( | ||
| inferenceProfileIdentifier: string, | ||
| awsRegion: string, | ||
| awsAccessKeyId: string, | ||
| awsSecretAccessKey: string, | ||
| awsSessionToken?: string | ||
| ) => { | ||
| const url = `https://bedrock.${awsRegion}.amazonaws.com/inference-profiles/${encodeURIComponent(decodeURIComponent(inferenceProfileIdentifier))}`; |
There was a problem hiding this comment.
🔒 Security Issue Fix
Issue: The getInferenceProfile function doesn't validate the inferenceProfileIdentifier before using it in the URL, which could potentially lead to URL manipulation issues.
Fix: Add validation to ensure the inferenceProfileIdentifier is a valid ARN format before using it.
Impact: Prevents potential security issues related to URL manipulation.
| export const getInferenceProfile = async ( | |
| inferenceProfileIdentifier: string, | |
| awsRegion: string, | |
| awsAccessKeyId: string, | |
| awsSecretAccessKey: string, | |
| awsSessionToken?: string | |
| ) => { | |
| const url = `https://bedrock.${awsRegion}.amazonaws.com/inference-profiles/${encodeURIComponent(decodeURIComponent(inferenceProfileIdentifier))}`; | |
| export const getInferenceProfile = async ( | |
| inferenceProfileIdentifier: string, | |
| awsRegion: string, | |
| awsAccessKeyId: string, | |
| awsSecretAccessKey: string, | |
| awsSessionToken?: string | |
| ) => { | |
| if (!inferenceProfileIdentifier || !inferenceProfileIdentifier.startsWith('arn:aws')) { | |
| throw new Error('Invalid inference profile identifier format'); | |
| } | |
| const url = `https://bedrock.${awsRegion}.amazonaws.com/inference-profiles/${encodeURIComponent(decodeURIComponent(inferenceProfileIdentifier))}`; |
|
Important PR Review SkippedPR review skipped as per the configuration setting. Run a manually review by commenting /matter review 💡Tips to use Matter AICommand List
|
narengogi
changed the title
narengogi
changed the title
remove redundant cache write and check if cache function is available before invoking it
|
Important PR Review SkippedPR review skipped as per the configuration setting. Run a manually review by commenting /matter review 💡Tips to use Matter AICommand List
|
|
Important PR Review SkippedPR review skipped as per the configuration setting. Run a manually review by commenting /matter review 💡Tips to use Matter AICommand List
|
|
Important PR Review SkippedPR review skipped as per the configuration setting. Run a manually review by commenting /matter review 💡Tips to use Matter AICommand List
|
more on bedrock inference profiles:
https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles.html
Testing done:
arn:aws:bedrock:us-east-1:517194595696:application-inference-profile/s529qz7ddy06(both URI encoded and as a regular string), verified that cost calculation is working as intendedanthropic.claude-3-haiku-20240307-v1:0Guide
aws bedrock create-inference-profile \ --inference-profile-name inference-profile-test \ --model-source copyFrom=arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-v2:1Note
inference profiles are immutable
arn:aws:bedrock:us-east-1:517194595696:application-inference-profile/s529qz7ddy06)snippets for testing