fix: use IAM roles for bedrock access by nrfulton · Pull Request #578 · generative-computing/mellea

nrfulton · 2026-03-05T00:04:24Z

Misc PR

Type of PR

Bug Fix
New Feature
Documentation
Other

Description

Link to Issue: Fixes 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 · 2026-03-05T00:04:35Z

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

mergify · 2026-03-05T00:04:59Z

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 · 2026-03-05T00:11:27Z

@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 · 2026-03-05T02:37:18Z

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 · 2026-03-05T02:39:30Z

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 · 2026-03-12T21:31:21Z

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:

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
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
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 · 2026-03-13T13:00:18Z

@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 · 2026-03-13T14:48:36Z

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 · 2026-03-13T19:24:53Z

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 · 2026-05-07T21:46:15Z

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:

logprobs will be an ad hoc field (see discussion on feat: bedrock aws access key #849)
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 · 2026-05-07T22:06:29Z

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!

jakelorocco · 2026-05-13T17:16:05Z

    hf_model_name="openai/gpt-oss-20b",  # OpenAI GPT-OSS 20B
    ollama_name="gpt-oss:20b",  # Ollama
    bedrock_name="openai.gpt-oss-20b",
+    bedrock_litellm_name="bedrock/converse/openai.gpt-oss-120b-1:0",


Change to 20B.

jakelorocco · 2026-05-13T17:20:47Z

+def _assert_region(region: str | None) -> None:
+    resolved_region = (
+        region
+        or os.environ.get("AWS_REGION_NAME")
+        or os.environ.get("AWS_DEFAULT_REGION")
+        or os.environ.get("AWS_REGION")
+    )
+    assert resolved_region is not None, (
+        "you must specify a region: pass `region` explicitly or set AWS_REGION_NAME, AWS_DEFAULT_REGION, or AWS_REGION."
+    )


Do the OpenAI and LiteLLM SDKs check for these same env vars?

jakelorocco · 2026-05-13T17:22:24Z

+    # TODO litellm doesn't even appear to use this...?
+    backend._base_url = None  # type: ignore


Can we remove this? or add a comment why setting the _base_url to None is necessary?

jakelorocco · 2026-05-13T17:24:08Z


-def create_bedrock_mantle_backend(
+def create_bedrock_litellm_backend(
+    model_id: ModelIdentifier | str, region: str | None = None, num_retries: int = 15


15 seems like a high default for retries. I did a quick search and it seems like Litellm uses exponential waits. I couldn't find a number, but that seems like this might resolve in long hangs.

jakelorocco · 2026-05-13T17:24:42Z

+        # Extract top-level Ollama params that must not be forwarded into `options`.
+        logprobs = model_opts.pop("logprobs", None)
+        top_logprobs = model_opts.pop("top_logprobs", None)
+        MelleaLogger.get_logger().info(f"Tools for call: {tools.keys()}")


This should be moved back to the indented block above.

jakelorocco · 2026-05-13T17:27:04Z


 m = MelleaSession(
-    backend=create_bedrock_mantle_backend(
+    backend=create_bedrock_openai_backend(


This should be the litellm path given the import above.

planetf1

A few things separate from the existing thread. Inline comments above cover a misnamed function in the troubleshooting section, a boto3 import regression for base-install users, assert-as-config-validation (which gets stripped under -O), and the missing deprecation alias.

One more that cannot be pinned inline — docs/docs/integrations/bedrock.md line 45 still has a stray prose reference to create_bedrock_mantle_backend that was not updated by the rename.

psschwei · 2026-06-04T22:20:51Z

To try to help move this along, I pushed a commit that attempted to address the outstanding review items from @jakelorocco and @planetf1 . If we're not happy with it, I don't mind reverting either.

jakelorocco

Paul's changes addressed my concerns.

psschwei · 2026-06-17T15:32:03Z

+<!-- leothomas, please delete this line -->
+


@leothomas can you suggest a change (through the UI) to delete these two lines? the UI should now ask you to sign-off, so that should let us satisfy the bot and get this in

Hey @psschwei let me know if that worked as expected. Thank you!

psschwei · 2026-06-22T15:58:00Z

gentle ping @leothomas

leothomas · 2026-06-24T18:56:51Z

Hi @psschwei! I submitted a review with the requested change to remove the line you had indicated last week, please let me know if that's what you were looking for or if I've misunderstood the ask

psschwei · 2026-06-24T20:53:06Z

Weird, I don't see your review @leothomas

I reached out on another medium with an idea on how to get this resolved.

leothomas · 2026-06-17T16:19:27Z

+<!-- leothomas, please delete this line -->
+


Hey @psschwei let me know if that worked as expected. Thank you!

leothomas · 2026-06-24T21:38:03Z

Oh man @psschwei that's completely my fault - I hadn't submitted the review, so it was still in a pending state. I've just submitted it - let me know if that works

psschwei · 2026-06-24T21:55:32Z

DCO check passes now:

though technically it didn't include you as signing off :/ ... can you reply with a "I sign off on my commits on this PR"? then I'll brute force it

Adds IAM auth support to Bedrock, logprobs and retries support, boto3 auth verification, and renames create_bedrock_mantle_backend to create_bedrock_openai_backend. Includes related test and docs updates. Signed-off-by: Paul S. Schweigert <paul@paulschweigert.com> Co-authored-by: Nathan Fulton <gitcommit@nfulton.org> Signed-off-by: Nathan Fulton <gitcommit@nfulton.org> Co-authored-by: Leo Thomas <leothomas@users.noreply.github.com> Signed-off-by: Leo Thomas <leothomas@users.noreply.github.com>

Signed-off-by: Paul S. Schweigert <paul@paulschweigert.com>

Co-authored-by: Leo Thomas <leothomas@users.noreply.github.com> Signed-off-by: Paul Schweigert <paul@paulschweigert.com>

psschwei · 2026-06-24T22:12:37Z

So @leothomas if you can confirm signing off, we can merge like this:

leothomas · 2026-06-25T01:26:43Z

I sign off on my commits on this PR

leothomas · 2026-06-25T01:27:12Z

Big thanks for seeing this through @psschwei - apologies for all the complications with the PR merging process!

psschwei · 2026-06-25T03:37:33Z

psschwei · 2026-06-25T03:39:16Z

Big thanks for seeing this through @psschwei - apologies for all the complications with the PR merging process!

no worries! thanks for your patience across this process.

@jakelorocco would you mind taking one more quick look? I had to rebase to pick up some changes, which required a bit of reworking, so would be good to give it a quick once over before we merge

jakelorocco · 2026-06-25T12:34:07Z

I think we missed this last time; so feel free to open a follow-up PR / issue since this PR doesn't make it worse, but the region parameter for stringify_mantle_model_ids should probably be passed here?

was an easy fix, just did it here

jakelorocco · 2026-06-25T12:42:04Z

+    model_options: dict = {"num_retries": num_retries}
+    resolved_region = _resolve_region(region)
+    if resolved_region is not None:
+        model_options["aws_region_name"] = resolved_region
+
+    return LiteLLMBackend(
+        model_id=model_name, model_options=model_options, num_retries=num_retries
+    )


num_retries is getting passed in as both a model_option and directly as a parameter. We should just choose one way of doing this. I feel like as a model_option makes sense.

Signed-off-by: Paul S. Schweigert <paul@paulschweigert.com>

Signed-off-by: Jake LoRocco <jake.lorocco@ibm.com>

jakelorocco

I pushed a fix to remove num_retries completely from the litellm backend. The bedrock function was already passing it as a model option after @psschwei's changes. There's no need to have the field.

psschwei · 2026-06-25T20:20:57Z

alright, let's do this!

nrfulton mentioned this pull request Mar 5, 2026

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

Closed

nrfulton changed the title ~~Fix/bedrock aws access key~~ fix: use IAM roles for bedrock access Mar 5, 2026

github-actions Bot added the bug Something isn't working label May 7, 2026

nrfulton mentioned this pull request May 7, 2026

Fixes pre-commit check fails. leothomas/mellea#1

Closed

nrfulton marked this pull request as ready for review May 7, 2026 21:38

nrfulton requested a review from a team as a code owner May 7, 2026 21:38

nrfulton requested review from ajbozarth and planetf1 May 7, 2026 21:38

nrfulton mentioned this pull request May 7, 2026

feat: bedrock aws access key #849

Closed

7 tasks

nrfulton requested review from jakelorocco and removed request for planetf1 May 7, 2026 21:42

jakelorocco requested changes May 13, 2026

View reviewed changes