docs: Bedrock example. (#410)

* Adds Bedrock example.

* Adds README for bedrock example, and checks that bearer token is defined.

* removes bare except.

* Adds an openai example and moves boto3 to [litellm], removing [aws]

* feat: add bedrock_name to model identifiers and bedrock examples

* Adds pytest skip reasons.

* reverts change to .pre-commit-config.

* Revert "reverts change to .pre-commit-config."

This reverts commit 649743a.

* bedrock work.

* Fixes model ids.

* saving wokr on openai bedrock example.

* fixes pre-commit hooks.

* Some cleanup.

Author: nrfulton

docs/examples/bedrock/README.md

@@ -0,0 +1,74 @@
+# Using Mellea with Bedrock
+
+Mellea can be used with Bedrock models via Mellea's LiteLLM or OpenAI backend support.
+
+## Pre-requisites
+
+To get started:
+
+1. Set the `AWS_BEARER_TOKEN_BEDROCK` environment variable.
+
+```bash
+export AWS_BEARER_TOKEN_BEDROCK=<your API key goes here>
+```
+
+2. If you want to use litellm, you need to install the optional `litellm` dependencies:
+
+```python
+uv pip install mellea[litellm]
+```
+
+3. You can now use Bedrock. We've included some built-in model ids for convienance:
+
+
+```python
+from mellea import MelleaSession
+from mellea.backends.bedrock import create_bedrock_mantle_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")
+
+m = MelleaSession(backend=bedrock_oai_backend, ctx=ChatContext())
+
+print(m.chat("Tell me 3 facts about Amazon.").content)
+```
+
+You can also use your own model IDs as strings, as long as they're accessible using the [mantle endpoints](https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.html):
+
+```python
+from mellea import MelleaSession
+from mellea.backends.bedrock import create_bedrock_mantle_backend
+from mellea.stdlib.context import ChatContext
+
+bedrock_oai_backend = create_bedrock_mantle_backend(
+    model_id="qwen.qwen3-coder-480b-a35b-instruct", 
+    region="us-east-1"
+)
+
+m = MelleaSession(backend=bedrock_oai_backend, ctx=ChatContext())
+
+print(m.chat("Tell me 3 facts about Amazon.").content)
+```
+
+You can get a list of all models that are available at the mantle endpoint for a region by running this utility script:
+
+```python
+from mellea.backends.bedrock import stringify_mantle_model_ids
+REGION = "us-east-1" # change this to see other region availability.
+print(f"Available Models in {REGION}:\n{stringify_mantle_model_ids(region=REGION)}")
+```
+
+## Other Examples
+
+Using LiteLLM with Bedrock (based on [litellm's docs](https://docs.litellm.ai/docs/providers/bedrock)):
+
+```
+uv run bedrock_litellm_example.py
+```
+
+Stand-alone email using the mantle OpenAI backends:
+
+```python
+uv run bedrock_openai_example.py
+```

docs/examples/bedrock/bedrock_litellm_example.py

@@ -0,0 +1,35 @@
+# pytest: skip
+# SKIP REASON: Requires an AWS bearer token for Bedrock.
+#!/usr/bin/env -S uv run --script
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "mellea[litellm]",
+#   "boto3" # including so that this example works before the next release.
+# ]
+# ///
+import os
+
+import mellea
+
+try:
+    import boto3
+except Exception:
+    raise Exception(
+        "Using Bedrock requires separately installing boto3. "
+        "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 = "bedrock/converse/us.amazon.nova-pro-v1:0"
+
+m = mellea.start_session(backend_name="litellm", model_id=MODEL_ID)
+
+result = m.chat("Give me three facts about Amazon.")
+
+print(result.content)

docs/examples/bedrock/bedrock_openai_example.py

@@ -0,0 +1,31 @@
+# pytest: skip
+# SKIP REASON: Requires an AWS bearer token for Bedrock.
+#!/usr/bin/env -S uv run --script
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "mellea[openai]"
+# ]
+# ///
+import os
+
+from mellea import MelleaSession
+from mellea.backends import model_ids
+from mellea.backends.bedrock import create_bedrock_mantle_backend
+from mellea.backends.openai import OpenAIBackend
+from mellea.stdlib.context import ChatContext
+
+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>"
+)
+
+m = MelleaSession(
+    backend=create_bedrock_mantle_backend(model_id=model_ids.OPENAI_GPT_OSS_120B),
+    ctx=ChatContext(),
+)
+
+result = m.chat("Give me three facts about Amazon.")
+
+print(result.content)

mellea/backends/bedrock.py

@@ -0,0 +1,78 @@
+"""Helpers for creating bedrock backends from openai/litellm."""
+
+import os
+
+from openai import OpenAI
+from openai.pagination import SyncPage
+
+from mellea.backends.model_ids import ModelIdentifier
+from mellea.backends.openai import OpenAIBackend
+
+
+def _make_region_for_uri(region: str | None):
+    if region is None:
+        region = "us-east-1"
+    return region
+
+
+def _make_mantle_uri(region: str | None = None):
+    region_for_uri = _make_region_for_uri(region)
+    uri = f"https://bedrock-mantle.{region_for_uri}.api.aws/v1"
+    return uri
+
+
+def list_mantle_models(region: str | None = None) -> list:
+    """Helper function get getting all models available at a bedrock-mantle endpoint."""
+    uri = _make_mantle_uri(region)
+    client = OpenAI(base_url=uri, api_key=os.environ["AWS_BEARER_TOKEN_BEDROCK"])
+    ms = client.models.list()
+    all_models = list(ms)
+    assert ms.next_page_info() is None
+    return all_models
+
+
+def stringify_mantle_model_ids(region: str | None = None) -> str:
+    """Helper function for getting a human-readable list of all models available at the mantle endpoint for an AWS region."""
+    models = list_mantle_models()
+    model_names = "\n * ".join([str(m.id) for m in models])
+    return f" * {model_names}"
+
+
+def create_bedrock_mantle_backend(
+    model_id: ModelIdentifier | str, region: str | None = None
+) -> OpenAIBackend:
+    """Returns an OpenAI backend that points to Bedrock mantle for model `model_id`."""
+    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"
+            )
+        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 != ""
+
+    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"
+    )
+
+    uri = _make_mantle_uri(region=region)
+
+    models = list_mantle_models(region)
+    if model_name not in [m.id for m in models]:
+        raise Exception(
+            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?"
+        )
+
+    backend = OpenAIBackend(
+        model_id=model_name,  # sic: do not pass the model_id itself!
+        base_url=uri,
+        api_key=os.environ["AWS_BEARER_TOKEN_BEDROCK"],
+    )
+    return backend

mellea/backends/model_ids.py

@@ -17,6 +17,7 @@ class ModelIdentifier:
     watsonx_name: str | None = None
     mlx_name: str | None = None
     openai_name: str | None = None
+    bedrock_name: str | None = None
 
     hf_tokenizer_name: str | None = None  # if None, is the same as hf_model_name
 
@@ -154,10 +155,12 @@ class ModelIdentifier:
 OPENAI_GPT_OSS_20B = ModelIdentifier(
     hf_model_name="openai/gpt-oss-20b",  # OpenAI GPT-OSS 20B
     ollama_name="gpt-oss:20b",  # Ollama
+    bedrock_name="openai.gpt-oss-20b",
 )
 OPENAI_GPT_OSS_120B = ModelIdentifier(
     hf_model_name="openai/gpt-oss-120b",  # OpenAI GPT-OSS 120B
     ollama_name="gpt-oss:120b",  # Ollama
+    bedrock_name="openai.gpt-oss-120b",
 )
 
 ###########################

pyproject.toml

@@ -89,7 +89,8 @@ vllm = [
 ]
 
 litellm = [
-    "litellm>=1.76"
+    "litellm>=1.76",
+    "boto3", # for Bedrock support.
 ]
 
 watsonx = [

test/backends/test_bedrock.py

@@ -0,0 +1,53 @@
+import os
+
+import openai
+import pytest
+
+import mellea.backends.model_ids
+import mellea.backends.model_ids as model_ids
+from mellea import MelleaSession
+from mellea.backends.bedrock import create_bedrock_mantle_backend
+from mellea.backends.openai import OpenAIBackend
+from mellea.stdlib.context import ChatContext
+
+# Skip entire module in CI since the single test is qualitative
+pytestmark = [
+    pytest.mark.llm,
+    pytest.mark.skipif(
+        "AWS_BEARER_TOKEN_BEDROCK" not in os.environ.keys(),
+        reason="Skipping Bedrock backend tests if $AWS_BEARER_TOKEN_BEDROCK is not set.",
+    ),
+]
+
+
+def _is_bedrock_model(model_id: model_ids.ModelIdentifier):
+    return model_id.bedrock_name is not None
+
+
+def test_model_ids_exist():
+    bedrock_models = [
+        getattr(mellea.backends.model_ids, name)
+        for name in dir(mellea.backends.model_ids)
+        if "bedrock_name" in dir(getattr(mellea.backends.model_ids, name))
+        and _is_bedrock_model(getattr(mellea.backends.model_ids, name))
+    ]
+
+    # non_bedrock_models = [
+    #     getattr(mellea.backends.model_ids, name)
+    #     for name in dir(mellea.backends.model_ids)
+    #     if "bedrock_name" not in dir(getattr(mellea.backends.model_ids, name))
+    #     and "ollama_name" in dir(getattr(mellea.backends.model_ids, name))
+    # ]
+
+    print(f"Found {len(bedrock_models)} bedrock-supported models.")
+    for model in bedrock_models:
+        print(f"Checking {model.bedrock_name}")
+        m = MelleaSession(
+            backend=create_bedrock_mantle_backend(model_id=model), ctx=ChatContext()
+        )
+        print(m.chat("What is 1+1?").content)
+
+
+if __name__ == "__main__":
+    test_model_ids_exist()
+    # pytest.main([__file__])