Skip to content

fix: use IAM roles for bedrock access #578

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

Merged
psschwei merged 5 commits into from
Jun 25, 2026
Merged

fix: use IAM roles for bedrock access #578

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
eeac9d8
fix: bedrock aws_access_key_id and related improvements
psschwei Jun 17, 2026
30ba9d2
workaround to get a dco signoff
psschwei Jun 17, 2026
2f9b6f0
Update README.md
psschwei Jun 24, 2026
db3bc9b
review fixes
psschwei Jun 25, 2026
8177743
fix: remove num_retries from litellm; use model_opts
jakelorocco Jun 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

# Mellea — build predictable AI without guesswork


Inside every AI-powered pipeline, the unreliable part is the same: the LLM call itself.
Silent failures, untestable outputs, no guarantees.
Mellea is a Python library for writing *generative programs* — replacing brittle prompts and flaky agents
Expand Down
22 changes: 11 additions & 11 deletions docs/docs/integrations/bedrock.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,18 +21,18 @@ Export it before running Mellea:
export AWS_BEARER_TOKEN_BEDROCK=your-bedrock-key
```

## Connecting with `create_bedrock_mantle_backend`
## Connecting with `create_bedrock_openai_backend`

```python
# Requires: mellea
# Returns: ModelOutputThunk
from mellea import MelleaSession
from mellea.backends import model_ids
from mellea.backends.bedrock import create_bedrock_mantle_backend
from mellea.backends.bedrock import create_bedrock_openai_backend
from mellea.stdlib.context import ChatContext

m = MelleaSession(
backend=create_bedrock_mantle_backend(model_id=model_ids.OPENAI_GPT_OSS_120B),
backend=create_bedrock_openai_backend(model_id=model_ids.OPENAI_GPT_OSS_120B),
ctx=ChatContext(),
)

Expand All @@ -41,7 +41,7 @@ print(str(result))
# Output will vary — LLM responses depend on model and temperature.
```

`create_bedrock_mantle_backend` returns an [`OpenAIBackend`](../reference/glossary#backend) pointed at the Bedrock
`create_bedrock_openai_backend` returns an [`OpenAIBackend`](../reference/glossary#backend) pointed at the Bedrock
Mantle endpoint. Pass it to [`MelleaSession`](../reference/glossary#melleasession) as shown above. It reads `AWS_BEARER_TOKEN_BEDROCK` from the environment and checks
that the requested model is available in the target region before returning.

Expand All @@ -53,10 +53,10 @@ The default region is `us-east-1`. Pass `region` to target a different region:
# Requires: mellea
# Returns: MelleaSession
from mellea import MelleaSession
from mellea.backends.bedrock import create_bedrock_mantle_backend
from mellea.backends.bedrock import create_bedrock_openai_backend

m = MelleaSession(
backend=create_bedrock_mantle_backend(
backend=create_bedrock_openai_backend(
model_id="amazon.nova-pro-v1:0",
region="eu-west-1",
)
Expand All @@ -72,11 +72,11 @@ model ID string directly:
# Requires: mellea
# Returns: MelleaSession
from mellea import MelleaSession
from mellea.backends.bedrock import create_bedrock_mantle_backend
from mellea.backends.bedrock import create_bedrock_litellm_backend

m = MelleaSession(
backend=create_bedrock_mantle_backend(
model_id="anthropic.claude-3-haiku-20240307-v1:0"
backend=create_bedrock_litellm_backend(
model_id="bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0"
)
)
```
Expand Down Expand Up @@ -126,7 +126,7 @@ available model IDs and credential setup.
**`AWS_BEARER_TOKEN_BEDROCK` not set:**

```text
AssertionError: Using AWS Bedrock requires setting a AWS_BEARER_TOKEN_BEDROCK environment variable.
RuntimeError: Using AWS Bedrock requires setting a AWS_BEARER_TOKEN_BEDROCK environment variable.
```

Export the environment variable before running your script:
Expand All @@ -143,7 +143,7 @@ Model X is not supported in region us-east-1.

Either enable model access for the requested model in your AWS account at
[Bedrock Model Access](https://us-east-1.console.aws.amazon.com/bedrock/home#/model-access),
or pass a different `region` to `create_bedrock_mantle_backend`.
or pass a different `region` to `create_bedrock_openai_backend`.

Comment thread
planetf1 marked this conversation as resolved.
## Vision support

Expand Down
8 changes: 4 additions & 4 deletions docs/examples/bedrock/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,11 +23,11 @@ uv pip install mellea[litellm]

```python
from mellea import MelleaSession
from mellea.backends.bedrock import create_bedrock_mantle_backend
from mellea.backends.bedrock import create_bedrock_openai_backend
from mellea.backends.model_ids import OPENAI_GPT_OSS_120B
from mellea.stdlib.context import ChatContext

bedrock_oai_backend = create_bedrock_mantle_backend(model_id=OPENAI_GPT_OSS_120B, region="us-east-1")
bedrock_oai_backend = create_bedrock_openai_backend(model_id=OPENAI_GPT_OSS_120B, region="us-east-1")

m = MelleaSession(backend=bedrock_oai_backend, ctx=ChatContext())

Expand All @@ -38,10 +38,10 @@ You can also use your own model IDs as strings, as long as they're accessible us

```python
from mellea import MelleaSession
from mellea.backends.bedrock import create_bedrock_mantle_backend
from mellea.backends.bedrock import create_bedrock_openai_backend
from mellea.stdlib.context import ChatContext

bedrock_oai_backend = create_bedrock_mantle_backend(
bedrock_oai_backend = create_bedrock_openai_backend(
model_id="qwen.qwen3-coder-480b-a35b-instruct",
region="us-east-1"
)
Expand Down
17 changes: 8 additions & 9 deletions docs/examples/bedrock/bedrock_litellm_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,9 @@
import os

import mellea
from mellea.backends.bedrock import create_bedrock_litellm_backend
from mellea.backends.model_ids import MISTRALAI_DEVSTRAL_2_123B
from mellea.stdlib.context import SimpleContext

try:
import boto3
Expand All @@ -20,16 +23,12 @@
"Run `uv pip install mellea[litellm]`"
)

assert "AWS_BEARER_TOKEN_BEDROCK" in os.environ.keys(), (
"Using AWS Bedrock requires setting a AWS_BEARER_TOKEN_BEDROCK environment variable. "
"Generate a key from the AWS console at: https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/api-keys?tab=long-term "
"Then run `export AWS_BEARER_TOKEN_BEDROCK=<insert your key here>"
)
MODEL_ID = MISTRALAI_DEVSTRAL_2_123B

MODEL_ID = "bedrock/converse/us.amazon.nova-pro-v1:0"
backend = create_bedrock_litellm_backend(MODEL_ID)
ctx = SimpleContext()
m = mellea.MelleaSession(backend, ctx)

m = mellea.start_session(backend_name="litellm", model_id=MODEL_ID)

result = m.chat("Give me three facts about Amazon.")
result = m.chat("What model am I talking to rn?")

print(result.content)
7 changes: 4 additions & 3 deletions docs/examples/bedrock/bedrock_openai_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@

from mellea import MelleaSession
from mellea.backends import model_ids
from mellea.backends.bedrock import create_bedrock_mantle_backend
from mellea.backends.bedrock import create_bedrock_openai_backend
from mellea.backends.model_ids import OPENAI_GPT_OSS_120B
from mellea.backends.openai import OpenAIBackend
from mellea.stdlib.context import ChatContext

Expand All @@ -22,10 +23,10 @@
)

m = MelleaSession(
backend=create_bedrock_mantle_backend(model_id=model_ids.OPENAI_GPT_OSS_120B),
backend=create_bedrock_openai_backend(model_id=OPENAI_GPT_OSS_120B),
ctx=ChatContext(),
)

result = m.chat("Give me three facts about Amazon.")
result = m.chat("What model am I talking to rn?")

print(result.content)
171 changes: 154 additions & 17 deletions mellea/backends/bedrock.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,80 @@
"""Helpers for creating bedrock backends from openai/litellm."""

import logging
import os
import warnings

from openai import OpenAI
from openai.pagination import SyncPage

from mellea.backends.litellm import LiteLLMBackend
from mellea.backends.model_ids import ModelIdentifier
from mellea.backends.openai import OpenAIBackend


def _resolve_region(region: str | None) -> str | None:
return (
region
or os.environ.get("AWS_REGION_NAME")
or os.environ.get("AWS_DEFAULT_REGION")
or os.environ.get("AWS_REGION")
)


def _assert_region(region: str | None) -> None:
if _resolve_region(region) is None:
raise ValueError(
"you must specify a region: pass `region` explicitly or set AWS_REGION_NAME, AWS_DEFAULT_REGION, or AWS_REGION."
)


def _assert_bedrock_auth() -> None:
Comment thread
jakelorocco marked this conversation as resolved.
"""Raises if no valid AWS credentials can be resolved.

Accepts any credential source that boto3 supports:
- Static env vars (AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY)
- Named profile (AWS_PROFILE or ~/.aws/credentials)
- ECS task role (AWS_CONTAINER_CREDENTIALS_RELATIVE_URI)
- EC2 / ECS instance profile (IMDSv2)
- LiteLLM-specific Bedrock API key (AWS_BEARER_TOKEN_BEDROCK)

Raises:
ImportError: If boto3 is not installed (install via the `litellm` extra).
RuntimeError: If no AWS credentials can be resolved.
"""
if "AWS_BEARER_TOKEN_BEDROCK" in os.environ:
return

try:
import boto3
import botocore.exceptions
except ImportError as e:
raise ImportError(
"boto3 is required to validate AWS credentials. "
"Please `pip install mellea[litellm]` (which includes boto3) "
"or set AWS_BEARER_TOKEN_BEDROCK to skip credential validation."
) from e

# botocore logs a credential-resolution message on every boto3.Session() call. Suppress it.
logging.getLogger("botocore.credentials").setLevel(logging.WARNING)

try:
creds = boto3.Session().get_credentials()
if creds is None:
raise botocore.exceptions.NoCredentialsError()
# Resolve to catch expired/invalid assume-role chains early.
creds.get_frozen_credentials()
except botocore.exceptions.NoCredentialsError:
raise RuntimeError(
"No AWS credentials found. Provide one of:\n"
" - AWS_BEARER_TOKEN_BEDROCK (Bedrock API key)\n"
" - AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY\n"
" - AWS_PROFILE pointing to a configured profile\n"
" - An IAM role attached to the instance/task (EC2, ECS, Lambda)"
)
except botocore.exceptions.NoRegionError:
pass # Credentials exist; region is validated separately.


def _make_region_for_uri(region: str | None):
if region is None:
region = "us-east-1"
Expand Down Expand Up @@ -48,12 +114,63 @@ def stringify_mantle_model_ids(region: str | None = None) -> str:
Returns:
Newline-separated string of model IDs prefixed with `" * "`.
"""
models = list_mantle_models()
models = list_mantle_models(region)
model_names = "\n * ".join([str(m.id) for m in models])
return f" * {model_names}"


def create_bedrock_mantle_backend(
def create_bedrock_litellm_backend(
model_id: ModelIdentifier | str, region: str | None = None, num_retries: int = 3
) -> LiteLLMBackend:
"""Returns a LiteLLM backend that points to Bedrock for model `model_id`.

Use this instead of `create_bedrock_openai_backend` when you need auth with an AWS_ACCESS_KEY_ID.

Args:
model_id: A `ModelIdentifier` (must have `bedrock_litellm_name`) or a raw
litellm-format Bedrock model ID string (e.g. `"bedrock/..."`).
region: AWS region. If `None`, falls back to AWS_REGION_NAME /
AWS_DEFAULT_REGION / AWS_REGION env vars.
num_retries: Retry budget for LiteLLM. LiteLLM uses exponential backoff,
so keep this low to avoid long hangs on persistent failures.

Raises:
ValueError: If no region can be resolved or `model_id` does not specify a
bedrock litellm name.
RuntimeError: If no AWS credentials can be resolved.
"""
_assert_bedrock_auth()
_assert_region(region)
Comment thread
jakelorocco marked this conversation as resolved.

model_name = ""
match model_id:
case ModelIdentifier():
if model_id.bedrock_litellm_name is None:
raise ValueError(
f"We do not have a known bedrock model identifier for {model_id}. If Bedrock supports this model, please pass the model_id string directly and open an issue to add the model id: https://github.com/generative-computing/mellea/issues/new"
)
else:
model_name = model_id.bedrock_litellm_name
case str():
model_name = model_id
if model_name == "":
raise ValueError(
f"Model identifier {model_id} does not specify a bedrock_name."
)

# Pass num_retries and the resolved region through model_options so litellm
# picks them up. num_retries goes here rather than as a direct LiteLLMBackend
# argument so there is a single source of truth; passing it both ways would
# send it to litellm.acompletion twice.
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)


def create_bedrock_openai_backend(
model_id: ModelIdentifier | str, region: str | None = None
Comment thread
planetf1 marked this conversation as resolved.
) -> OpenAIBackend:
"""Return an OpenAI backend that points to Bedrock mantle for the given model.
Expand All @@ -70,37 +187,40 @@ def create_bedrock_mantle_backend(
via AWS Bedrock Mantle.

Raises:
Exception: If `model_id` is a `ModelIdentifier` with no `bedrock_name`
set.
AssertionError: If the `AWS_BEARER_TOKEN_BEDROCK` environment variable is
ValueError: If `model_id` is a `ModelIdentifier` with no `bedrock_name`
set, or if the specified model is not available in the target region.
RuntimeError: If the `AWS_BEARER_TOKEN_BEDROCK` environment variable is
not set.
Exception: If the specified model is not available in the target region.
"""
model_name = ""
match model_id:
case ModelIdentifier() if model_id.bedrock_name is None:
raise Exception(
f"We do not have a known bedrock model identifier for {model_id}. If Bedrock supports this model, please pass the model_id string directly and open an issue to add the model id: https://github.com/generative-computing/mellea/issues/new"
raise ValueError(
f"We do not have a known bedrock model identifier for {model_id}. If Bedrock supports this model, please pass the model_id string directly and open an issue to add the model id: https://github.com/generative-computing/mellea/issues/new"
)
case ModelIdentifier() if model_id.bedrock_name is not None:
assert model_id.bedrock_name is not None # for type checker help.
model_name = model_id.bedrock_name
case str():
model_name = model_id
assert model_name != ""
if model_name == "":
raise ValueError(
f"Model identifier {model_id} does not specify a bedrock_name."
)

assert "AWS_BEARER_TOKEN_BEDROCK" in os.environ.keys(), (
"Using AWS Bedrock requires setting a AWS_BEARER_TOKEN_BEDROCK environment variable.\n\nTo proceed:\n"
"\n\t1. Generate a key from the AWS console at: https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/api-keys?tab=long-term "
"\n\t2. Run `export AWS_BEARER_TOKEN_BEDROCK=<insert your key here>\n"
"If you need to use normal AWS credentials instead of a bedrock-specific bearer token, please open an issue: https://github.com/generative-computing/mellea/issues/new"
)
if "AWS_BEARER_TOKEN_BEDROCK" not in os.environ:
raise RuntimeError(
"Using AWS Bedrock requires setting a AWS_BEARER_TOKEN_BEDROCK environment variable.\n\nTo proceed:\n"
"\n\t1. Generate a key from the AWS console at: https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/api-keys?tab=long-term "
"\n\t2. Run `export AWS_BEARER_TOKEN_BEDROCK=<insert your key here>`\n"
"If you need to use normal AWS credentials instead of a bedrock-specific bearer token, use create_bedrock_litellm_backend instead."
)

uri = _make_mantle_uri(region=region)

models = list_mantle_models(region)
if model_name not in [m.id for m in models]:
raise Exception(
raise ValueError(
f"Model {model_name} is not supported in region {_make_region_for_uri(region=region)}.\nSupported models are:\n{stringify_mantle_model_ids(region)}\n\nPerhaps change regions or check that model access for {model_name} is not gated on Bedrock?"
)

Expand All @@ -110,3 +230,20 @@ def create_bedrock_mantle_backend(
api_key=os.environ["AWS_BEARER_TOKEN_BEDROCK"],
)
return backend


def create_bedrock_mantle_backend(
model_id: ModelIdentifier | str, region: str | None = None
) -> OpenAIBackend:
"""Deprecated alias for `create_bedrock_openai_backend`.

.. deprecated::
Use `create_bedrock_openai_backend` instead. This shim will be removed
in a future release.
"""
warnings.warn(
"create_bedrock_mantle_backend is deprecated; use create_bedrock_openai_backend instead.",
DeprecationWarning,
stacklevel=2,
)
return create_bedrock_openai_backend(model_id=model_id, region=region)
Loading
Loading