From eeac9d8cadd92e4b95f425c459706804f881598d Mon Sep 17 00:00:00 2001 From: "Paul S. Schweigert" Date: Wed, 17 Jun 2026 11:10:55 -0400 Subject: [PATCH 1/5] fix: bedrock aws_access_key_id and related improvements 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 Co-authored-by: Nathan Fulton Signed-off-by: Nathan Fulton Co-authored-by: Leo Thomas Signed-off-by: Leo Thomas --- docs/docs/integrations/bedrock.md | 22 +-- docs/examples/bedrock/README.md | 8 +- .../bedrock/bedrock_litellm_example.py | 17 +- .../bedrock/bedrock_openai_example.py | 7 +- mellea/backends/bedrock.py | 169 ++++++++++++++++-- mellea/backends/litellm.py | 38 +++- mellea/backends/model_ids.py | 7 + mellea/backends/ollama.py | 5 + mellea/plugins/policies.py | 4 +- test/backends/test_bedrock.py | 39 +++- ...bedrock_unit.py => test_bedrock_openai.py} | 25 ++- 11 files changed, 270 insertions(+), 71 deletions(-) rename test/backends/{test_bedrock_unit.py => test_bedrock_openai.py} (77%) diff --git a/docs/docs/integrations/bedrock.md b/docs/docs/integrations/bedrock.md index 65b62bcbf..74c1168fa 100644 --- a/docs/docs/integrations/bedrock.md +++ b/docs/docs/integrations/bedrock.md @@ -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(), ) @@ -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. @@ -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", ) @@ -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" ) ) ``` @@ -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: @@ -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`. ## Vision support diff --git a/docs/examples/bedrock/README.md b/docs/examples/bedrock/README.md index a995c1dd5..964bfc613 100644 --- a/docs/examples/bedrock/README.md +++ b/docs/examples/bedrock/README.md @@ -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()) @@ -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" ) diff --git a/docs/examples/bedrock/bedrock_litellm_example.py b/docs/examples/bedrock/bedrock_litellm_example.py index e7812a123..e74f42062 100644 --- a/docs/examples/bedrock/bedrock_litellm_example.py +++ b/docs/examples/bedrock/bedrock_litellm_example.py @@ -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 @@ -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=" -) +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) diff --git a/docs/examples/bedrock/bedrock_openai_example.py b/docs/examples/bedrock/bedrock_openai_example.py index 5f20525a8..eef798dab 100644 --- a/docs/examples/bedrock/bedrock_openai_example.py +++ b/docs/examples/bedrock/bedrock_openai_example.py @@ -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 @@ -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) diff --git a/mellea/backends/bedrock.py b/mellea/backends/bedrock.py index 500d5d4a0..4f930faad 100644 --- a/mellea/backends/bedrock.py +++ b/mellea/backends/bedrock.py @@ -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: + """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" @@ -53,7 +119,58 @@ def stringify_mantle_model_ids(region: str | None = None) -> str: 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) + + 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 the resolved region through model_options so litellm picks it up even + # when `region` was supplied explicitly rather than via env vars. + 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 + ) + + +def create_bedrock_openai_backend( model_id: ModelIdentifier | str, region: str | None = None ) -> OpenAIBackend: """Return an OpenAI backend that points to Bedrock mantle for the given model. @@ -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=\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=`\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?" ) @@ -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) diff --git a/mellea/backends/litellm.py b/mellea/backends/litellm.py index 0f0a81b7d..9da04543d 100644 --- a/mellea/backends/litellm.py +++ b/mellea/backends/litellm.py @@ -17,6 +17,7 @@ 'Please install them with: pip install "mellea[litellm]"' ) from e + from ..backends import model_ids from ..core import ( BaseModelSubclass, @@ -86,6 +87,7 @@ def __init__( formatter: ChatFormatter | None = None, base_url: str | None = None, model_options: dict | None = None, + num_retries: int = 0, ): """Initialize a LiteLLM-compatible backend for the given model ID and endpoint.""" super().__init__( @@ -110,6 +112,8 @@ def __init__( base_url if base_url is not None else "http://localhost:11434/v1" ) + self._num_retries = num_retries + # A mapping of common options for this backend mapped to their Mellea ModelOptions equivalent. # These are usually values that must be extracted before hand or that are common among backend providers. # OpenAI has some deprecated parameters. Those map to the same mellea parameter, but @@ -263,9 +267,17 @@ def _make_backend_specific_and_remove( standard_openai_subset = litellm.get_standard_openai_params(backend_specific) unknown_keys = [] # Keys that are unknown to litellm. unsupported_openai_params = [] # OpenAI params that are known to litellm but not supported for this model/provider. + # Bedrock-specific pass-through params that LiteLLM accepts but doesn't list as supported OpenAI params. + known_provider_passthrough = { + "additional_model_request_fields", + "additional_model_response_field_paths", + } + for key in backend_specific.keys(): if key not in supported_params: - if key in standard_openai_subset: + if key in known_provider_passthrough: + pass # Expected provider-specific params; no warning needed. + elif key in standard_openai_subset: # LiteLLM is pretty confident that this standard OpenAI parameter won't work. unsupported_openai_params.append(key) else: @@ -290,8 +302,9 @@ async def _generate_from_chat_context_standard( action: Component[C] | CBlock, ctx: Context, *, - _format: type[BaseModelSubclass] - | None = None, # Type[BaseModelSubclass] is a class object of a subclass of BaseModel + _format: ( + type[BaseModelSubclass] | None + ) = None, # Type[BaseModelSubclass] is a class object of a subclass of BaseModel model_options: dict | None = None, tool_calls: bool = False, ) -> ModelOutputThunk[C]: @@ -414,6 +427,7 @@ async def _generate_from_chat_context_standard( tools=formatted_tools, api_base=resolved_api_base, drop_params=True, # See note in `_make_backend_specific_and_remove`. + num_retries=self._num_retries, **extra_params, **reasoning_params, # type: ignore **model_specific_options, @@ -504,6 +518,16 @@ async def processing( if content_chunk is not None: mot._underlying_value += content_chunk + if getattr(choice, "logprobs", None) is not None: + mot._meta["logprobs"] = choice.logprobs + + # In some cases (converse API) Bedrock returns logprobs via additionalModelResponseFields. + additional_fields = getattr(chunk, "model_extra", {}) or {} + if "additionalModelResponseFields" in additional_fields: + mot._meta["additionalModelResponseFields"] = additional_fields[ + "additionalModelResponseFields" + ] + # Store the full response (includes usage) as a dict. mot.raw.response = chunk.model_dump() @@ -521,6 +545,12 @@ async def processing( if content_chunk is not None: mot._underlying_value += content_chunk + stream_logprobs = getattr(chunk.choices[0], "logprobs", None) + if stream_logprobs is not None: + if "logprobs" not in mot._meta: + mot._meta["logprobs"] = [] + mot._meta["logprobs"].append(stream_logprobs) + if mot.raw.streamed_chunks is None: mot.raw.streamed_chunks = [] mot.raw.streamed_chunks.append(chunk.choices[0].model_dump()) @@ -603,6 +633,7 @@ async def post_processing( } generate_log.action = mot._action generate_log.result = mot + mot._generate_log = generate_log # Non-streaming carries usage on the response; streaming already set it. @@ -697,6 +728,7 @@ async def _generate_from_raw( completion_response = await litellm.atext_completion( model=self._model_id, prompt=prompts, + num_retries=self._num_retries, api_base=user_api_base_raw or (self._base_url if self._explicit_base_url else None), **model_specific_options, diff --git a/mellea/backends/model_ids.py b/mellea/backends/model_ids.py index 7277b2afe..ffffd8cdb 100644 --- a/mellea/backends/model_ids.py +++ b/mellea/backends/model_ids.py @@ -35,6 +35,7 @@ class ModelIdentifier: mlx_name: str | None = None openai_name: str | None = None bedrock_name: str | None = None + bedrock_litellm_name: str | None = None hf_tokenizer_name: str | None = None # if None, is the same as hf_model_name context_length: int | None = None @@ -241,6 +242,10 @@ class ModelIdentifier: context_length=131072, ) +MISTRALAI_DEVSTRAL_2_123B = ModelIdentifier( + bedrock_name="mistral.devstral-2-123b", + bedrock_litellm_name="bedrock/converse/mistral.devstral-2-123b", +) ##################### #### Qwen models #### @@ -278,12 +283,14 @@ class ModelIdentifier: 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-20b-1:0", context_length=131072, ) 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", + bedrock_litellm_name="bedrock/converse/openai.gpt-oss-120b-1:0", context_length=131072, ) diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index 6f1266cec..b9c078f98 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -445,6 +445,9 @@ async def generate_from_chat_context( # they overwrite conflicting names. add_tools_from_context_actions(tools, [action]) MelleaLogger.get_logger().info(f"Tools for call: {tools.keys()}") + # 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) # Generate a chat response from ollama, using the chat messages. Can be either type since stream is passed as a model option. chat_response: Coroutine[ @@ -457,6 +460,8 @@ async def generate_from_chat_context( stream=model_opts.get(ModelOption.STREAM, False), options=self._make_backend_specific_and_remove(model_opts), format=_format.model_json_schema() if _format is not None else None, # type: ignore + logprobs=logprobs, + top_logprobs=top_logprobs, ) # type: ignore output = ModelOutputThunk(None) diff --git a/mellea/plugins/policies.py b/mellea/plugins/policies.py index b5ebb9e5c..2a9f16bbe 100644 --- a/mellea/plugins/policies.py +++ b/mellea/plugins/policies.py @@ -5,8 +5,8 @@ from typing import Any try: - from cpex.framework.hooks.policies import ( - HookPayloadPolicy, # type: ignore[import-not-found] + from cpex.framework.hooks.policies import ( # type: ignore[import-not-found] + HookPayloadPolicy, ) _HAS_PLUGIN_FRAMEWORK = True diff --git a/test/backends/test_bedrock.py b/test/backends/test_bedrock.py index 9b416f30c..f92ac06ec 100644 --- a/test/backends/test_bedrock.py +++ b/test/backends/test_bedrock.py @@ -1,10 +1,15 @@ +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.bedrock import ( + create_bedrock_litellm_backend, + create_bedrock_openai_backend, +) from mellea.backends.openai import OpenAIBackend from mellea.stdlib.context import ChatContext from test.predicates import require_api_key @@ -28,22 +33,38 @@ def test_model_ids_exist(): 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() + backend=create_bedrock_openai_backend(model_id=model), ctx=ChatContext() ) print(m.chat("What is 1+1?").content) +@pytest.mark.qualitative +@pytest.mark.skipif( + not os.environ.get("AWS_REGION_NAME") + and not os.environ.get("AWS_DEFAULT_REGION") + and not os.environ.get("AWS_REGION"), + reason="No AWS region set; cannot exercise the litellm bedrock path.", +) +def test_litellm_bedrock_chat(): + """Smoke test for the litellm bedrock path. + + Skipped unless an AWS_BEARER_TOKEN_BEDROCK is present (module-level skip) + and an AWS region is resolvable. Uses a string model_id so the test does not + depend on `model_ids.OPENAI_GPT_OSS_*.bedrock_litellm_name` staying stable. + """ + backend = create_bedrock_litellm_backend( + model_id="bedrock/converse/openai.gpt-oss-20b-1:0" + ) + m = MelleaSession(backend=backend, ctx=ChatContext()) + response = m.chat("Reply with the single word: ok") + assert response.content is not None + assert len(str(response.content).strip()) > 0 + + if __name__ == "__main__": test_model_ids_exist() # pytest.main([__file__]) diff --git a/test/backends/test_bedrock_unit.py b/test/backends/test_bedrock_openai.py similarity index 77% rename from test/backends/test_bedrock_unit.py rename to test/backends/test_bedrock_openai.py index 34d499847..76321cacd 100644 --- a/test/backends/test_bedrock_unit.py +++ b/test/backends/test_bedrock_openai.py @@ -13,7 +13,7 @@ from mellea.backends.bedrock import ( _make_mantle_uri, _make_region_for_uri, - create_bedrock_mantle_backend, + create_bedrock_openai_backend, ) from mellea.backends.model_ids import ModelIdentifier from mellea.backends.openai import OpenAIBackend @@ -48,49 +48,46 @@ def test_mantle_uri_custom_region(): assert uri == "https://bedrock-mantle.ap-northeast-1.api.aws/v1" -# --- create_bedrock_mantle_backend error paths --- +# --- create_bedrock_openai_backend error paths --- def test_model_identifier_without_bedrock_name_raises(): mid = ModelIdentifier(hf_model_name="some/model") with pytest.raises(Exception, match="do not have a known bedrock model identifier"): - create_bedrock_mantle_backend(mid) + create_bedrock_openai_backend(mid) def test_missing_env_var_raises(monkeypatch): monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("AWS_ACCESS_KEY_ID", raising=False) mid = ModelIdentifier(bedrock_name="some.model-id") - with pytest.raises(AssertionError, match="AWS_BEARER_TOKEN_BEDROCK"): - create_bedrock_mantle_backend(mid) + with pytest.raises(RuntimeError, match="AWS_BEARER_TOKEN_BEDROCK"): + create_bedrock_openai_backend(mid) -# --- create_bedrock_mantle_backend happy paths (mock network boundary) --- +# --- create_bedrock_openai_backend happy paths (mock network boundary) --- @patch("mellea.backends.bedrock.list_mantle_models", return_value=_FAKE_MODELS) def test_create_backend_with_model_identifier(mock_list, monkeypatch): monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "fake-token") mid = ModelIdentifier(bedrock_name="granite-3.3-8b") - backend = create_bedrock_mantle_backend(mid, region="eu-west-1") + backend = create_bedrock_openai_backend(mid, region="eu-west-1") assert isinstance(backend, OpenAIBackend) assert backend.model_id == "granite-3.3-8b" - assert backend._base_url == "https://bedrock-mantle.eu-west-1.api.aws/v1" mock_list.assert_called_once_with("eu-west-1") @patch("mellea.backends.bedrock.list_mantle_models", return_value=_FAKE_MODELS) def test_create_backend_with_string_model_id(mock_list, monkeypatch): monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "fake-token") - backend = create_bedrock_mantle_backend("llama-4-scout") - - assert isinstance(backend, OpenAIBackend) - assert backend.model_id == "llama-4-scout" - assert backend._base_url == "https://bedrock-mantle.us-east-1.api.aws/v1" + backend = create_bedrock_openai_backend("llama-4-scout") + assert backend.model_id == "llama-4-scout", "model id should be llama-4-scout" @patch("mellea.backends.bedrock.list_mantle_models", return_value=_FAKE_MODELS) def test_model_not_in_region_raises(mock_list, monkeypatch): monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "fake-token") with pytest.raises(Exception, match="not supported in region"): - create_bedrock_mantle_backend("nonexistent-model") + create_bedrock_openai_backend("nonexistent-model") From 30ba9d2418047fc3a75d2783c35cff5e77a644f1 Mon Sep 17 00:00:00 2001 From: "Paul S. Schweigert" Date: Wed, 17 Jun 2026 11:14:13 -0400 Subject: [PATCH 2/5] workaround to get a dco signoff Signed-off-by: Paul S. Schweigert --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index c7d955088..f0a82c45e 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,8 @@ # 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 From 2f9b6f01458f212f0012828e002d1d44d48bd637 Mon Sep 17 00:00:00 2001 From: Paul Schweigert Date: Wed, 24 Jun 2026 17:52:28 -0400 Subject: [PATCH 3/5] Update README.md Co-authored-by: Leo Thomas Signed-off-by: Paul Schweigert --- README.md | 1 - 1 file changed, 1 deletion(-) diff --git a/README.md b/README.md index f0a82c45e..bf445dc52 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,6 @@ # 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. From db3bc9b1c1bfbbb26600a5c64c516f066f95204e Mon Sep 17 00:00:00 2001 From: "Paul S. Schweigert" Date: Thu, 25 Jun 2026 12:15:31 -0400 Subject: [PATCH 4/5] review fixes Signed-off-by: Paul S. Schweigert --- mellea/backends/bedrock.py | 12 ++++---- test/backends/test_bedrock_openai.py | 43 ++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/mellea/backends/bedrock.py b/mellea/backends/bedrock.py index 4f930faad..65ff810cd 100644 --- a/mellea/backends/bedrock.py +++ b/mellea/backends/bedrock.py @@ -114,7 +114,7 @@ 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}" @@ -158,16 +158,16 @@ def create_bedrock_litellm_backend( f"Model identifier {model_id} does not specify a bedrock_name." ) - # Pass the resolved region through model_options so litellm picks it up even - # when `region` was supplied explicitly rather than via env vars. + # 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, num_retries=num_retries - ) + return LiteLLMBackend(model_id=model_name, model_options=model_options) def create_bedrock_openai_backend( diff --git a/test/backends/test_bedrock_openai.py b/test/backends/test_bedrock_openai.py index 76321cacd..0e395042c 100644 --- a/test/backends/test_bedrock_openai.py +++ b/test/backends/test_bedrock_openai.py @@ -13,8 +13,11 @@ from mellea.backends.bedrock import ( _make_mantle_uri, _make_region_for_uri, + create_bedrock_litellm_backend, create_bedrock_openai_backend, + stringify_mantle_model_ids, ) +from mellea.backends.litellm import LiteLLMBackend from mellea.backends.model_ids import ModelIdentifier from mellea.backends.openai import OpenAIBackend @@ -91,3 +94,43 @@ def test_model_not_in_region_raises(mock_list, monkeypatch): monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "fake-token") with pytest.raises(Exception, match="not supported in region"): create_bedrock_openai_backend("nonexistent-model") + + +# --- stringify_mantle_model_ids --- + + +@patch("mellea.backends.bedrock.list_mantle_models", return_value=_FAKE_MODELS) +def test_stringify_mantle_model_ids_passes_region(mock_list): + stringify_mantle_model_ids("eu-west-1") + mock_list.assert_called_once_with("eu-west-1") + + +# --- create_bedrock_litellm_backend --- + + +@patch("mellea.backends.bedrock._assert_bedrock_auth") +def test_litellm_backend_num_retries_single_source(mock_auth, monkeypatch): + """num_retries is wired through model_options only, not as a direct arg. + + Passing it both ways would forward num_retries to litellm.acompletion twice + (once via the constructor param, once via model_specific_options), raising a + TypeError at generation time. Keeping it solely in model_options is the + single source of truth. + """ + monkeypatch.setenv("AWS_REGION", "us-east-1") + backend = create_bedrock_litellm_backend( + "bedrock/converse/openai.gpt-oss-20b-1:0", num_retries=7 + ) + + assert isinstance(backend, LiteLLMBackend) + assert backend.model_options["num_retries"] == 7 + # The direct LiteLLMBackend num_retries param keeps its default (0); the + # value flows exclusively through model_options. + assert backend._num_retries == 0 + + +@patch("mellea.backends.bedrock._assert_bedrock_auth") +def test_litellm_backend_resolved_region_in_model_options(mock_auth, monkeypatch): + monkeypatch.setenv("AWS_REGION", "ap-northeast-1") + backend = create_bedrock_litellm_backend("bedrock/converse/openai.gpt-oss-20b-1:0") + assert backend.model_options["aws_region_name"] == "ap-northeast-1" From 81777431d3bccdf50769d2667b0eda2f0c1855f7 Mon Sep 17 00:00:00 2001 From: Jake LoRocco Date: Thu, 25 Jun 2026 15:24:47 -0400 Subject: [PATCH 5/5] fix: remove num_retries from litellm; use model_opts Signed-off-by: Jake LoRocco --- mellea/backends/litellm.py | 5 ----- test/backends/test_bedrock_openai.py | 3 --- 2 files changed, 8 deletions(-) diff --git a/mellea/backends/litellm.py b/mellea/backends/litellm.py index 9da04543d..daf1c2126 100644 --- a/mellea/backends/litellm.py +++ b/mellea/backends/litellm.py @@ -87,7 +87,6 @@ def __init__( formatter: ChatFormatter | None = None, base_url: str | None = None, model_options: dict | None = None, - num_retries: int = 0, ): """Initialize a LiteLLM-compatible backend for the given model ID and endpoint.""" super().__init__( @@ -112,8 +111,6 @@ def __init__( base_url if base_url is not None else "http://localhost:11434/v1" ) - self._num_retries = num_retries - # A mapping of common options for this backend mapped to their Mellea ModelOptions equivalent. # These are usually values that must be extracted before hand or that are common among backend providers. # OpenAI has some deprecated parameters. Those map to the same mellea parameter, but @@ -427,7 +424,6 @@ async def _generate_from_chat_context_standard( tools=formatted_tools, api_base=resolved_api_base, drop_params=True, # See note in `_make_backend_specific_and_remove`. - num_retries=self._num_retries, **extra_params, **reasoning_params, # type: ignore **model_specific_options, @@ -728,7 +724,6 @@ async def _generate_from_raw( completion_response = await litellm.atext_completion( model=self._model_id, prompt=prompts, - num_retries=self._num_retries, api_base=user_api_base_raw or (self._base_url if self._explicit_base_url else None), **model_specific_options, diff --git a/test/backends/test_bedrock_openai.py b/test/backends/test_bedrock_openai.py index 0e395042c..8ebfc1a64 100644 --- a/test/backends/test_bedrock_openai.py +++ b/test/backends/test_bedrock_openai.py @@ -124,9 +124,6 @@ def test_litellm_backend_num_retries_single_source(mock_auth, monkeypatch): assert isinstance(backend, LiteLLMBackend) assert backend.model_options["num_retries"] == 7 - # The direct LiteLLMBackend num_retries param keeps its default (0); the - # value flows exclusively through model_options. - assert backend._num_retries == 0 @patch("mellea.backends.bedrock._assert_bedrock_auth")