Skip to content

feat: Add prompt routing to AmazonBedrockChatGenerator #2220

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 11 commits into
base: main
Choose a base branch
from
Open

feat: Add prompt routing to AmazonBedrockChatGenerator #2220

wants to merge 11 commits into from

Conversation

Amnah199
Copy link
Contributor

@Amnah199 Amnah199 commented Aug 25, 2025 •
edited
Loading

Related Issues

Proposed Changes:

  • Allow passing prompt_router_arn during initialization
  • Make model param optional (shouldn't be a breaking change)
  • Add an example

How did you test it?

  • Update unit tests

Notes for the reviewer

Checklist

@github-actions github-actions bot added integration:amazon-bedrock type:documentation Improvements or additions to documentation labels Aug 25, 2025
@Amnah199 Amnah199 marked this pull request as ready for review August 25, 2025 14:45
@Amnah199 Amnah199 requested a review from a team as a code owner August 25, 2025 14:45
@Amnah199 Amnah199 requested review from sjrl and removed request for a team August 25, 2025 14:45
@sjrl sjrl requested a review from medsriha August 25, 2025 15:28
@sjrl
Copy link
Contributor

sjrl commented Aug 25, 2025

Hey @medsriha it would be great if you could check if this solution works for your use case!

@sjrl
Copy link
Contributor

sjrl commented Aug 25, 2025

@Amnah199 Taking a look at the code you added it seems like we are enabling two features here.

  1. Allow a user to create a new PromptRouter in Bedrock by passing a prompt_router_config
  2. Allow a user to use an existing PromptRouter config in Bedrock by basically finding an already existing promptRouterArn. You do this currently by doing a lookup and comparing the name provided in prompt_router_config.

Looking at the issue originally opened by @medsriha it seems like we really only need to support the second case currently. Basically allow a user to pass a promptRouterArn as input. I'm not sure if we need the fully prompt router creation to be added currently. WDYT?

sjrl
sjrl reviewed Aug 25, 2025
View reviewed changes
integrations/amazon_bedrock/tests/test_chat_generator.py Outdated
Comment on lines 147 to 151
def test_to_dict_with_prompt_router_config(self, mock_boto3_session, boto3_config):
"""
Test that the to_dict method returns the correct dictionary without aws credentials
"""
generator = AmazonBedrockChatGenerator(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In addition to this test it would be good to add an integration test where we actually run the generator with the prompt router and check that it returns the expected output.

@Amnah199
Copy link
Contributor Author

Amnah199 commented Aug 25, 2025 •
edited
Loading

@sjrl I see what you mean. For simplicity, if we just allow passing router ARNs than the user can simply pass the ARNs for default routers or the ones they have configured in the AWS console.

On the other hand, in the context issue we also have:

Bedrock prompt routers allow users to define custom routing logic.

If we want to support this aspect within our component, we can keep my implementation. I am open to both depending on what are the cons of keeping this support for router_config.

@sjrl
Copy link
Contributor

sjrl commented Aug 25, 2025

@sjrl I see what you mean. For simplicity, if we just allow passing router ARNs than the user can simply pass the ARNs for default routers or the ones they have configured in the AWS console.

On the other hand, in the context issue we also have:

Bedrock prompt routers allow users to define custom routing logic.

If we want to support this aspect within our component, we can keep my implementation. I am open to both depending on what are the cons of keeping this support for router_config.

@medsriha could you weigh-in and let us know if you need both of these features?

@medsriha
Copy link
Member

Thanks for pushing this forward, @Amnah199 and @sjrl. The ask was a bit ambiguous, I apologize. I think to keep it simple is to support reusing an existing prompt router via ARN for now. That way, users can create and manage their routers directly in AWS, while Haystack just consumes them. We can always revisit “create-on-the-fly” support later if there’s a demand.

Amnah199
Amnah199 commented Aug 26, 2025
View reviewed changes
...edrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py
Comment on lines +247 to +251
resolved_router_arn = resolve_secret(self.prompt_router_arn)
bedrock_client = session.client("bedrock", config=config)
prompt_router = bedrock_client.get_prompt_router(
promptRouterArn=resolved_router_arn,
)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@sjrl I believe we can show a specific error message in case of invalid ARN. If you agree I can raise an exception here with the message.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah that could be helpful. What did you have in mind?

sjrl
sjrl reviewed Aug 26, 2025
View reviewed changes
...edrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py
prompt_router = bedrock_client.get_prompt_router(
promptRouterArn=resolved_router_arn,
)
self.model = prompt_router["promptRouterArn"]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just to clarify, is this value prompt_router["promptRouterArn"] the same as the one provided by the user so prompt_router_arn at init time?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes it would be the same. Here we verify if it actually is a valid ARN.

Copy link
Contributor

@sjrl sjrl Aug 26, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay I see. Perhaps I could propose an alternative approach. If self.prompt_router_arn works in the same way as we treat self.model do we really need to make a new variable?

Would specifying an prompt_router_arn in the model field at init time already work with the existing integration?

sjrl
sjrl reviewed Aug 26, 2025
View reviewed changes
...edrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py
@@ -158,6 +163,8 @@ def __init__(
streaming_callback: Optional[StreamingCallbackT] = None,
boto3_config: Optional[Dict[str, Any]] = None,
tools: Optional[Union[List[Tool], Toolset]] = None,
*,
prompt_router_arn: Optional[Secret] = None,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just to double check, does prompt_router_arn need to be a Secret? What was the motivation to make it a secret?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The custom router ARN contains your region and AWS account ID, which probably we dont want to expose.

@Amnah199
Copy link
Contributor Author

@medsriha We realized that the existing generator implementation already supports passing a router ARN to the model parameter during initialization, so it should work as is.

The only potential reason to introduce a separate prompt_router_arn as a Secret param would be to avoid exposing the AWS account ID (since it’s embedded in custom router ARNs). If that’s not an immediate concern, I’d suggest testing your routers with the current setup first and letting us know how it goes.

@medsriha
Copy link
Member

@Amnah199 turns out your first solution was exactly what we needed. I had thought we weren’t supposed to set up a new ARN via Haystack, but what we actually need is a JSON configuration that tells Bedrock which model to use and when.
https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-routing.html#prompt-routing-use

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@sjrl sjrl sjrl left review comments

@medsriha medsriha Awaiting requested review from medsriha

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
integration:amazon-bedrock type:documentation Improvements or additions to documentation
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

Support AWS Bedrock Prompt Router Invocation in Generators
4 participants
@Amnah199 @sjrl @medsriha @anakin87