fix: use IAM roles for bedrock access

[nrfulton]

Misc PR

Type of PR

• Bug Fix
• New Feature
• Documentation
• Other

Description

• Link to Issue: Feat: Native AWS SigV4 (IAM) Support for Bedrock OpenAI-Compatible Endpoints #564

Testing

• Tests added to the respective file if code was changed
• New code has 100% coverage if code as added
• Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

[github-actions]

The PR description has been updated. Please fill out the template for your PR to be reviewed.

[mergify]

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

• title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

[nrfulton]

@jakelorocco you are tagged as a reviewer because I changed core.py to only call post-processing if there are no exceptions. The status quo continues to give us pain.

[jakelorocco]

I think that sounds good. I will review once the changes are done. We will need to make sure the finally block / telemetry stuff doesn't cause issues with this change.

[jakelorocco]

Looks like there's another draft PR open to more comprehensively fix the post process issue: #580 (comment). I will take a closer look tomorrow.

[leothomas]

Hey folks - thank you so much for adding support for the BedRock non-mantle endpoints through the LiteLLM backend! I'm currently working on integrating Mellea with a fork of IBM's FactReasoner. I made a couple additional (local) edits to this PR in order to get it to work for me, which include:

1. Validating authentication using boto3: this PR just checks for AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables but users can also be authenticated through named AWS profiles (AWS_PROFILE env var) or something like an EC2 instance or Lambda function would directly have an IAM role attached. So I update it to use boto3 to validate that it can reach the Bedrock service, regardless of how it's authneticated
2. Added explicit handling of logprobs: when executing queries against a Custom Model Import, with a model that supports logprob output (eg: Llama-X.X-instruct), I'm able to use the model_options parameter in the mfuncs.ainstruct function to request logprobs output alongside the main response. I also had to add logprobs support to LiteLLM, which I've forked and I'll be opening a PR for that
3. Added a num_retries parameter to the bedrock.create_beckrock_litellm_backend() factory to take advantage of liteLLM's native retry-with-backoff - this is important for Custom Model Imports as they are cycled off of Bedrock's compute infrastructure after a period of un-used and therefore have a bit of a coldstart that needs to be managed.

Let me know if these changes would be valuable to contribue back to the main Mellea repo, I'm happy to either open a PR from my own fork (which would be branched off of this PR), or wait till this PR is merged to open mine, or any other way to contribute the changes that makes sense with your workflow; I'm just trying to avoid forking on top of a fork! (the edits are all in the backend/bedrock.py and backends/litellm.py files)

Once again, thank for getting the LiteLLM Bedrock backend working!

[jakelorocco]

@leothomas, I think all these changes sound good. I will check in with @nrfulton today and see if we can get this base PR merged so that your PRs can be opened. I can't test this myself.

The logprobs is definitely something we've been looking into supporting for a little while (mainly just unifying it across backends). We'd love to see your implementation for litellm. Thank you!

[leothomas]

Hey there! I pulled @nrfulton 's branch into my own forks and then added the edits mentioned above. Here's the branch if you want to check it out: https://github.com/leothomas/mellea/tree/bedrock_aws_access_key_id

@jakelorocco I'll wait for y'all to merge this PR to open one from my branch!

[leothomas]

Also, here's the LiteLLM fork with LogProbs: https://github.com/leothomas/litellm/tree/feature/logprobs

I still have to write tests and all that good stuff before opening a PR against LiteLLM

[nrfulton]

I changed up the auto-selected reviewers to match the reviewers assigned to #849 .

FYI @ajbozarth and @jakelorocco: TL;DR: same state as #849. This is the OG PR for that work. Switched here because maintainer commit rights were disabled for the #849 branch and there were a number of failing pre-commit hooks + a merge required to proceed.

I think this is good to go if tests pass. Please give a once over since there's been a fair amount of changes since this PR was first opened.

The main design decisions:

1. logprobs will be an ad hoc field (see discussion on feat: bedrock aws access key #849)
2. the test should already be skipped if the AWS token isn't present. We don't have that token in the testing infra, so it'll be skipped in nightlies.

FYI, all tests pass on my local machine with AWS_BEARER_TOKEN_BEDROCK set.

$ export AWS_BEARER_TOKEN_BEDROCK=XXXXXXXXX; pytest test/backends/test_bedrock_openai.py
================================================================================== test session starts ==================================================================================
platform darwin -- Python 3.12.9, pytest-9.0.3, pluggy-1.6.0
rootdir: /Users/nathan/dev/mellea
configfile: pyproject.toml
plugins: nbmake-1.5.5, recording-0.13.4, Faker-40.15.0, cov-7.1.0, xdist-3.8.0, json-report-1.5.0, timeout-2.4.0, metadata-3.1.1, asyncio-1.3.0, langsmith-0.8.3, anyio-4.13.0
timeout: 900.0s
timeout method: signal
timeout func_only: False
asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 8 items                                                                                                                                                                       

test/backends/test_bedrock_openai.py ........                                                                                                                                     [100%]

==================================================================================== tests coverage =====================================================================================
___________________________________________________________________ coverage: platform darwin, python 3.12.9-final-0 ____________________________________________________________________

Coverage HTML written to dir htmlcov
Coverage JSON written to file coverage.json
=================================================================================== 8 passed in 5.91s ===================================================================================
$ export AWS_BEARER_TOKEN_BEDROCK=XXXXXXXXX; pytest test/backends/test_bedrock.py
================================================================================== test session starts ==================================================================================
platform darwin -- Python 3.12.9, pytest-9.0.3, pluggy-1.6.0
rootdir: /Users/nathan/dev/mellea
configfile: pyproject.toml
plugins: nbmake-1.5.5, recording-0.13.4, Faker-40.15.0, cov-7.1.0, xdist-3.8.0, json-report-1.5.0, timeout-2.4.0, metadata-3.1.1, asyncio-1.3.0, langsmith-0.8.3, anyio-4.13.0
timeout: 900.0s
timeout method: signal
timeout func_only: False
asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 1 item                                                                                                                                                                        

test/backends/test_bedrock.py .                                                                                                                                                   [100%]

==================================================================================== tests coverage =====================================================================================
___________________________________________________________________ coverage: platform darwin, python 3.12.9-final-0 ____________________________________________________________________

Coverage HTML written to dir htmlcov
Coverage JSON written to file coverage.json
================================================================================== 1 passed in 10.37s ===================================================================================

[leothomas]

Thanks for following up on this and getting it over the finish line @nrfulton - again please let me know if I can support in any way!