Add Amazon Bedrock Guardrails integration #2
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This commit adds the ability to use custom models imported into AWS Bedrock through the OpenAI-compatible API interface. Key features include: - User-friendly model IDs that include the model name (e.g., mistral-7b-instruct-id:custom.a1b2c3d4) - Support for both custom models and imported models - Detailed documentation in CUSTOM_MODELS_IMPLEMENTATION.md - Updates to README.md and Usage.md - Custom models are always enabled by default 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
This file provides guidance for Claude Code assistant when working with this repository, including commands and code style guidelines. Updated to use pipx for running ruff. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
- Update schema.py to handle user-friendly custom model IDs - Update model.py router to preserve user-friendly IDs in responses 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
- Add EcrAccountId parameter to both Lambda and Fargate templates - Default to the current official account ID (366590864501) - Use the parameter in ECR repository URLs and IAM policy This change allows users to easily deploy with custom ECR repositories. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Enhances custom model streaming support to handle: 1. Missing contentType in stream chunks 2. Support for 'generation' field format used by some models This ensures that custom imported models can be used with streaming mode.
Enhances custom model handling to support additional response formats: 1. Add support for 'generation' field in response 2. Add support for 'prompt_token_count' and 'generation_token_count' fields 3. Add debug logging of response keys to help troubleshoot This ensures that custom imported models can be used with both streaming and non-streaming modes.
These tests verify that: 1. Custom models appear in the models list API 2. Custom models can be invoked successfully with a non-empty response 3. Custom models support streaming mode The tests require at least one custom model to be set up in the AWS account.
This utility script allows direct testing of custom imported models using boto3: 1. Tests direct invocation with the AWS Bedrock runtime 2. Tests streaming invocation with the AWS Bedrock runtime 3. Parses different response formats used by custom models 4. Extracts complete texts from streaming chunks Useful for debugging or verifying custom model behavior outside of the gateway API.
- Extract helper function to reduce code duplication - Simplify custom model listing logic - Improve model lookup to handle friendly IDs and AWS IDs efficiently - Add helper methods to extract completion text and usage - Replace DEBUG blocks with logger.debug() calls - Streamline validation and response handling 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
- Add gateway-level guardrail configuration options - Apply guardrails to both standard and custom models - Support request-level guardrail parameters - Add documentation in chat.py example request 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
msharp9
approved these changes
Apr 25, 2025
There was a problem hiding this comment.
I was so confused at first since I was just looking for the guardrails stuff, but the majority of this PR is actually you adding the custom names for the custom models.
I am a bit suspect about how the custom naming is going to work with model versions and running into name conflicts there, but overall looks good.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
Test plan
🤖 Generated with Claude Code